diff --git a/.editorconfig b/.editorconfig index 88ff663335..1e61d25216 100644 --- a/.editorconfig +++ b/.editorconfig @@ -16,6 +16,9 @@ trim_trailing_whitespace = false [*.yml] indent_size = 2 +[*.yaml] +indent_size = 2 + [*.scala] indent_size = 2 diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index bfa317d57c..c1a86c9731 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -* @superhx @SCNieh @ShadowySpirits @Chillax-0v0 +* @superhx @Gezi-lzq @1sonofqiu @woshigaopp diff --git a/.github/ISSUE_TEMPLATE/bug-report.md b/.github/ISSUE_TEMPLATE/bug-report.md deleted file mode 100644 index 258e8e384a..0000000000 --- a/.github/ISSUE_TEMPLATE/bug-report.md +++ /dev/null @@ -1,41 +0,0 @@ ---- -name: "\U0001F41B Bug report" -about: Something is not working -title: "[BUG] " -labels: bug -assignees: '' - ---- - -### Version & Environment - - - - -### What went wrong? - - - -### What should have happened instead? - - - -### How to reproduce the issue? - -1. -2. -3. - -### Additional information - -Please attach any relevant logs, backtraces, or metric charts. diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml new file mode 100644 index 0000000000..fcf6ff0d19 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug-report.yml @@ -0,0 +1,105 @@ +name: "🐛 Bug Report" +description: Report a problem or unexpected behavior in AutoMQ +title: "[BUG] " +labels: [bug] +assignees: [your-github-username] +body: + - type: markdown + attributes: + value: | + Thanks for reporting a bug. Please fill out the form below to help us diagnose and fix the issue. + + - type: input + id: version + attributes: + label: AutoMQ Version + description: Run `automq --version` to find this. + placeholder: e.g., v0.1.0 + validations: + required: true + + - type: input + id: os + attributes: + label: Operating System + description: Use `uname` or check `/etc/os-release`. + placeholder: e.g., Ubuntu 20.04 + validations: + required: true + + - type: dropdown + id: install-method + attributes: + label: Installation Method + description: How did you install AutoMQ? + options: + - source + - binary + - docker + - package manager + - other + validations: + required: true + + - type: textarea + id: hardware + attributes: + label: Hardware Configuration + description: Describe CPU, memory, disk, and other relevant hardware specs. + placeholder: e.g., 4-core CPU, 16GB RAM, SSD + validations: + required: false + + - type: textarea + id: software + attributes: + label: Other Relevant Software + description: Include any additional tools or versions (e.g., Kafka version, monitoring tools). + placeholder: e.g., Kafka v3.6.0, Prometheus v2.45.0 + validations: + required: false + + - type: textarea + id: problem + attributes: + label: What Went Wrong? + description: Describe the unexpected behavior or error message. + placeholder: | + e.g. Broker crashes when receiving a large number of messages. + Error: "Connection refused: No available brokers" + validations: + required: true + + - type: textarea + id: expected + attributes: + label: What Should Have Happened Instead? + description: Describe what you expected to happen. + placeholder: e.g., The broker should have handled the load gracefully. + validations: + required: true + + - type: textarea + id: reproduction + attributes: + label: Steps to Reproduce + description: Provide step-by-step instructions to reproduce the issue. + placeholder: | + 1. Start AutoMQ broker with default config. + 2. Send a large number of messages. + 3. Observe the broker behavior. + validations: + required: true + + - type: textarea + id: additional + attributes: + label: Additional Information + description: Add logs, stack traces, metrics, configuration, or screenshots here. + placeholder: | + - Logs + - Backtraces + - Configuration files + - Screenshots + validations: + required: false diff --git a/.github/ISSUE_TEMPLATE/feature-request.md b/.github/ISSUE_TEMPLATE/feature-request.md deleted file mode 100644 index 425a34435a..0000000000 --- a/.github/ISSUE_TEMPLATE/feature-request.md +++ /dev/null @@ -1,16 +0,0 @@ ---- -name: "\U0001F680 Feature request" -about: Suggest an idea for AutoMQ -title: "[Enhancement]" -labels: enhancement -assignees: '' - ---- - -### Who is this for and what problem do they have today? - - -### Why is solving this problem impactful? - - -### Additional notes diff --git a/.github/ISSUE_TEMPLATE/feature-request.yml b/.github/ISSUE_TEMPLATE/feature-request.yml new file mode 100644 index 0000000000..4e3d9f8718 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature-request.yml @@ -0,0 +1,58 @@ +name: 🚀 Feature Request +description: Suggest a new idea or improvement for AutoMQ +title: "[Feature Request] " +labels: + - enhancement +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to suggest a feature! Please fill out the form below as clearly as possible. + + - type: input + id: target-audience + attributes: + label: Who is this feature for? + description: Who are the users that would benefit from this feature? + placeholder: e.g. Engineers needing real-time AutoMQ monitoring + validations: + required: true + + - type: textarea + id: current-problem + attributes: + label: What problem are they facing today? + description: Describe the issue or limitation they experience without this feature. + placeholder: e.g. No built-in dashboard to monitor metrics like throughput, latency, etc. + validations: + required: true + + - type: textarea + id: impact + attributes: + label: Why is solving this impactful? + description: Explain the importance of this feature and how it benefits users, the project, or the community. + placeholder: e.g. Improves system visibility, reduces debugging time, enables proactive alerting... + validations: + required: true + + - type: textarea + id: solution + attributes: + label: Proposed solution + description: Describe your proposed solution or idea in detail. + placeholder: | + - Web dashboard with real-time charts + - Accessible via the AutoMQ admin panel + - Uses minimal resources + validations: + required: true + + - type: textarea + id: notes + attributes: + label: Additional notes + description: Add any related issues, dependencies, or initial work done. + placeholder: e.g. Related to issue #123, might require metrics export setup + validations: + required: false diff --git a/.github/ISSUE_TEMPLATE/good-first-issue.md b/.github/ISSUE_TEMPLATE/good-first-issue.md deleted file mode 100644 index 336bf6349d..0000000000 --- a/.github/ISSUE_TEMPLATE/good-first-issue.md +++ /dev/null @@ -1,43 +0,0 @@ ---- -name: "⭐ Good first issue" -about: Design good first issue for new contributors -title: "[Good first issue] " -labels: good first issue -assignees: '' - ---- - -### Background - - - -### What's our expectation for the issue - - - -### How to started - - -### Reference -- [Kafka Official Document](https://kafka.apache.org/documentation/) diff --git a/.github/ISSUE_TEMPLATE/good-first-issue.yml b/.github/ISSUE_TEMPLATE/good-first-issue.yml new file mode 100644 index 0000000000..778600dba5 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/good-first-issue.yml @@ -0,0 +1,56 @@ +name: "⭐ Good First Issue" +description: Design and describe a task suitable for first-time contributors +title: "[Good First Issue] " +labels: [good first issue] +assignees: [your-github-username] +body: + - type: markdown + attributes: + value: | + Thank you for helping us make AutoMQ better for new contributors! Please fill out the details of this good first issue. + + - type: textarea + id: background + attributes: + label: Background + description: Explain why this issue is being created and any relevant context. + placeholder: | + e.g. Currently, Kafka's logs are stored only on the local disk. We want to also store them in cloud object storage like S3 to improve reliability and enable querying. + validations: + required: true + + - type: textarea + id: expectation + attributes: + label: What's Our Expectation for This Issue? + description: Describe what the desired outcome of this issue is. + placeholder: | + e.g. Log files should still exist locally, but also be uploaded to S3 using the following path format: s3://bucket-name/automq/cluster-id/broker-id/logs/xx + validations: + required: true + + - type: textarea + id: getting-started + attributes: + label: How to Get Started + description: Help contributors understand how to begin working on the issue. + placeholder: | + - Precondition: + - Understand how AutoMQ writes logs to the local filesystem. + - Main Classes: + - `LogWriter` + - `CloudUploader` + - Tips: + - Refer to the `LogWriter` class to understand the current logic. + validations: + required: true + + - type: textarea + id: references + attributes: + label: Reference Links + description: Include any helpful links, documents, or code references. + placeholder: | + e.g. [Kafka Official Documentation](https://kafka.apache.org/documentation/) + validations: + required: false diff --git a/.github/workflows/auto-assign-issue.yml b/.github/workflows/auto-assign-issue.yml new file mode 100644 index 0000000000..d2d1026fe6 --- /dev/null +++ b/.github/workflows/auto-assign-issue.yml @@ -0,0 +1,25 @@ +name: Auto Assign Issue +on: + issue_comment: + types: [created] + +jobs: + assign-issue: + runs-on: ubuntu-latest + steps: + - name: Check if comment contains '/assign' + if: contains(github.event.comment.body, '/assign') + uses: actions/github-script@v6 + with: + script: | + try { + await github.rest.issues.addAssignees({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.payload.issue.number, + assignees: [context.payload.comment.user.login] + }); + console.log(`Assigned issue #${context.payload.issue.number} to @${context.payload.comment.user.login}`); + } catch (error) { + console.error('Error assigning issue:', error); + } diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d6bb6926c4..ffbcb72626 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -16,7 +16,7 @@ name: Build on: pull_request: - types: [ opened, reopened, synchronize ] + types: [ opened, reopened, synchronize, ready_for_review ] push: branches: [ "main" ] diff --git a/.github/workflows/build_automq.yml b/.github/workflows/build_automq.yml index d33df8a6a6..e3ccaf5f68 100644 --- a/.github/workflows/build_automq.yml +++ b/.github/workflows/build_automq.yml @@ -49,7 +49,7 @@ jobs: - name: Setup Gradle uses: gradle/gradle-build-action@v2.9.0 - name: Checkstyle - run: ./gradlew --build-cache rat checkstyleMain checkstyleTest + run: ./gradlew --build-cache rat checkstyleMain checkstyleTest spotlessJavaCheck spotbugs: name: "Spotbugs" runs-on: ${{ matrix.os }} diff --git a/.github/workflows/docker-bitnami-release.yaml b/.github/workflows/docker-bitnami-release.yaml new file mode 100644 index 0000000000..7266cb0387 --- /dev/null +++ b/.github/workflows/docker-bitnami-release.yaml @@ -0,0 +1,67 @@ +name: Docker Bitnami Release + +on: + workflow_dispatch: + push: + tags: + - '[0-9]+.[0-9]+.[0-9]+' + - '[0-9]+.[0-9]+.[0-9]+-rc[0-9]+' + + +jobs: + docker-release: + name: Docker Image Release + strategy: + matrix: + platform: [ "ubuntu-24.04" ] + jdk: ["17"] + runs-on: ${{ matrix.platform }} + permissions: + contents: write + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Set up JDK ${{ matrix.jdk }} + uses: actions/setup-java@v3 + with: + java-version: ${{ matrix.jdk }} + distribution: "zulu" + - name: Setup Gradle + uses: gradle/gradle-build-action@v2.12.0 + - name: Get project version + id: get_project_version + run: | + project_version=$(./gradlew properties | grep "version:" | awk '{print $2}') + echo "PROJECT_VERSION=${project_version}" >> $GITHUB_OUTPUT + + - name: Build TarGz + run: | + ./gradlew -Pprefix=automq-${{ github.ref_name }}_ --build-cache --refresh-dependencies clean releaseTarGz + + # docker image release + - name: Cp TarGz to Docker Path + run: | + cp ./core/build/distributions/automq-${{ github.ref_name }}_kafka-${{ steps.get_project_version.outputs.PROJECT_VERSION }}.tgz ./container/bitnami + - name: Determine Image Tags + id: image_tags + run: | + echo "tags=${{ secrets.DOCKERHUB_USERNAME }}/automq:${{ github.ref_name }}-bitnami" >> $GITHUB_OUTPUT + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_READ_WRITE_TOKEN }} + - name: Build and push + uses: docker/build-push-action@v5 + with: + context: ./container/bitnami + push: true + tags: ${{ steps.image_tags.outputs.tags }} + platforms: linux/amd64,linux/arm64 diff --git a/.github/workflows/docker-kafka-release.yaml b/.github/workflows/docker-kafka-release.yaml new file mode 100644 index 0000000000..98ebd3e053 --- /dev/null +++ b/.github/workflows/docker-kafka-release.yaml @@ -0,0 +1,70 @@ +name: AutoMQ Kafka Docker Release + +on: + workflow_dispatch: + inputs: + tag: + description: 'AutoMQ Version Tag' + required: false + type: string + workflow_run: + workflows: ["GitHub Release"] + types: + - completed + +env: + KAFKA_VERSION: "3.9.0" + +jobs: + automq-kafka-release: + name: AutoMQ Kafka Docker Image Release + strategy: + matrix: + platform: [ "ubuntu-24.04" ] + jdk: [ "17" ] + runs-on: ${{ matrix.platform }} + permissions: + contents: write + steps: + - name: Checkout Code + uses: actions/checkout@v4 + + - name: Get release tag + run: | + if [[ "${{ github.event_name }}" == "workflow_dispatch" && -n "${{ github.event.inputs.tag }}" ]]; then + TAG="${{ github.event.inputs.tag }}" + # use the latest tag if not specified + elif [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then + TAG=$(git ls-remote --tags https://github.com/AutoMQ/automq.git | grep -v '\^{}' | tail -1 | sed 's/.*refs\/tags\///') + else + TAG="${{ github.event.workflow_run.head_branch }}" + fi + + AUTOMQ_URL="https://github.com/AutoMQ/automq/releases/download/${TAG}/automq-${TAG}_kafka-${KAFKA_VERSION}.tgz" + + { + echo "AUTOMQ_VERSION=${TAG}-kafka" + echo "AUTOMQ_URL=${AUTOMQ_URL}" + } >> $GITHUB_ENV + + - name: Set up Python 3.10 + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_READ_WRITE_TOKEN }} + + - name: Build and Push Docker Image + run: | + python3 -m venv .venv + source .venv/bin/activate + .venv/bin/pip install setuptools + + cd docker + python3 docker_release.py \ + ${{ secrets.DOCKERHUB_USERNAME }}/automq:${AUTOMQ_VERSION} \ + --kafka-url ${AUTOMQ_URL} diff --git a/.github/workflows/docker-release.yml b/.github/workflows/docker-release.yml index ca4dcac4a3..f57da03f9b 100644 --- a/.github/workflows/docker-release.yml +++ b/.github/workflows/docker-release.yml @@ -1,6 +1,7 @@ name: Docker Release on: + workflow_dispatch: push: tags: - '[0-9]+.[0-9]+.[0-9]+' @@ -12,7 +13,7 @@ jobs: name: Docker Image Release strategy: matrix: - platform: [ "ubuntu-22.04" ] + platform: [ "ubuntu-24.04" ] jdk: ["17"] runs-on: ${{ matrix.platform }} permissions: diff --git a/.github/workflows/docker-strimzi-release.yaml b/.github/workflows/docker-strimzi-release.yaml new file mode 100644 index 0000000000..6cfac8a239 --- /dev/null +++ b/.github/workflows/docker-strimzi-release.yaml @@ -0,0 +1,84 @@ +name: Docker Strimzi Release + +on: + workflow_dispatch: + inputs: + tag: + description: 'AutoMQ Version Tag' + required: false + type: string + workflow_run: + workflows: ["GitHub Release"] + types: + - completed + +env: + KAFKA_VERSION: "3.9.0" + STRIMZI_REPO: "https://github.com/AutoMQ/strimzi-kafka-operator.git" + STRIMZI_BRANCH: "main" + +jobs: + strimzi-release: + name: Strimzi Image Release + if: ${{ github.event.workflow_run.conclusion == 'success' || github.event_name == 'workflow_dispatch' }} + strategy: + matrix: + platform: [ "ubuntu-24.04" ] + jdk: ["17"] + runs-on: ${{ matrix.platform }} + permissions: + contents: write + steps: + - name: Checkout Code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Get release tag + run: | + if [[ "${{ github.event_name }}" == "workflow_dispatch" && -n "${{ github.event.inputs.tag }}" ]]; then + TAG="${{ github.event.inputs.tag }}" + # use the latest tag if not specified + elif [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then + TAG=$(git ls-remote --tags https://github.com/AutoMQ/automq.git | grep -v '\^{}' | tail -1 | sed 's/.*refs\/tags\///') + else + TAG="${{ github.event.workflow_run.head_branch }}" + fi + + AUTOMQ_URL="https://github.com/AutoMQ/automq/releases/download/${TAG}/automq-${TAG}_kafka-${KAFKA_VERSION}.tgz" + + { + echo "AUTOMQ_VERSION=${TAG}" + echo "AUTOMQ_URL=${AUTOMQ_URL}" + } >> $GITHUB_ENV + + - name: Set up JDK ${{ matrix.jdk }} + uses: actions/setup-java@v3 + with: + java-version: ${{ matrix.jdk }} + distribution: "zulu" + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_READ_WRITE_TOKEN }} + + - name: Build AutoMQ Strimzi Image + run: | + git clone --depth 1 --branch "${{ env.STRIMZI_BRANCH }}" "${{ env.STRIMZI_REPO }}" strimzi + cd strimzi + + chmod +x ./tools/automq/build-automq-image.sh + ./tools/automq/build-automq-image.sh \ + "${{ env.AUTOMQ_VERSION }}" \ + "${{ env.AUTOMQ_URL }}" \ + "${{ env.KAFKA_VERSION }}" \ + "${{ secrets.DOCKERHUB_USERNAME }}" \ + "automq" diff --git a/.github/workflows/e2e-run.yml b/.github/workflows/e2e-run.yml index 467840da08..b7d75346b4 100644 --- a/.github/workflows/e2e-run.yml +++ b/.github/workflows/e2e-run.yml @@ -57,12 +57,14 @@ jobs: run: ./tests/docker/run_tests.sh env: ESK_TEST_YML: ${{ inputs.test-yaml }} + _DUCKTAPE_OPTIONS: "--deflake 4" shell: bash - name: Run E2E tests with path if: ${{ inputs.test-path != '' }} run: ./tests/docker/run_tests.sh env: TC_PATHS: ${{ inputs.test-path }} + _DUCKTAPE_OPTIONS: "--deflake 4" shell: bash - name: Extract results id: extract-results diff --git a/.github/workflows/github-release.yml b/.github/workflows/github-release.yml index 7ba620fbfd..c7b128226e 100644 --- a/.github/workflows/github-release.yml +++ b/.github/workflows/github-release.yml @@ -30,32 +30,39 @@ jobs: uses: gradle/gradle-build-action@v2.12.0 - name: Build TarGz + id: build-targz run: | ./gradlew -Pprefix=automq-${{ github.ref_name }}_ --build-cache --refresh-dependencies clean releaseTarGz mkdir -p core/build/distributions/latest + LATEST_TAG=$(git tag --sort=-v:refname | grep -E '^[0-9]+\.[0-9]+\.[0-9]+$' | head -n 1) + echo "LATEST_TAG=$LATEST_TAG" + IS_LATEST="false" + if [ "$LATEST_TAG" == "${{ github.ref_name }}" ]; then + IS_LATEST=true + fi + echo "IS_LATEST=$IS_LATEST" >> $GITHUB_OUTPUT for file in core/build/distributions/automq-*.tgz; do if [[ ! "$file" =~ site-docs ]]; then - echo "Find latest tgz file: $file" - cp "$file" core/build/distributions/latest/automq-kafka-latest.tgz - break + if [ "$IS_LATEST" = "true" ]; then + echo "Find latest tgz file: $file" + cp "$file" core/build/distributions/latest/automq-kafka-latest.tgz + fi else echo "Skip and remove site-docs file: $file" rm "$file" fi done - - uses: jakejarvis/s3-sync-action@master - name: s3-upload-latest - if: ${{ github.repository_owner == 'AutoMQ' }} + - uses: tvrcgo/oss-action@master + name: upload-latest + if: ${{ github.repository_owner == 'AutoMQ' && steps.build-targz.outputs.IS_LATEST == 'true' }} with: - args: --follow-symlinks --delete - env: - AWS_S3_BUCKET: ${{ secrets.AWS_CN_PROD_BUCKET }} - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_CN_PROD_AK }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_CN_PROD_SK }} - AWS_REGION: 'cn-northwest-1' - SOURCE_DIR: 'core/build/distributions/latest' - DEST_DIR: 'community_edition/artifacts' + bucket: ${{ secrets.UPLOAD_BUCKET }} + key-id: ${{ secrets.UPLOAD_BUCKET_AK }} + key-secret: ${{ secrets.UPLOAD_BUCKET_SK }} + region: 'oss-cn-hangzhou' + assets: | + core/build/distributions/latest/automq-kafka-latest.tgz:community_edition/artifacts/automq-kafka-latest.tgz - name: GitHub Release uses: softprops/action-gh-release@v1 diff --git a/.github/workflows/nightly-extra-e2e.yml b/.github/workflows/nightly-extra-e2e.yml deleted file mode 100644 index 90177685a8..0000000000 --- a/.github/workflows/nightly-extra-e2e.yml +++ /dev/null @@ -1,61 +0,0 @@ -name: Nightly Extra E2E tests -on: - workflow_dispatch: - schedule: - - cron: '0 16 * * *' - -jobs: - benchmarks_e2e: - name: "Run benchmarks E2E Tests" - uses: ./.github/workflows/e2e-run.yml - if: ${{ github.repository_owner == 'AutoMQ' }} - with: - suite-id: "benchmarks" - test-path: "tests/kafkatest/benchmarks" - runner: "e2e" - connect_e2e_1: - name: "Run connect E2E Tests 1" - uses: ./.github/workflows/e2e-run.yml - if: ${{ github.repository_owner == 'AutoMQ' }} - with: - suite-id: "connect1" - test-yaml: "tests/suites/connect_test_suite1.yml" - runner: "e2e" - connect_e2e_2: - name: "Run connect E2E Tests 2" - uses: ./.github/workflows/e2e-run.yml - if: ${{ github.repository_owner == 'AutoMQ' }} - with: - suite-id: "connect2" - test-yaml: "tests/suites/connect_test_suite2.yml" - runner: "e2e" - connect_e2e_3: - name: "Run connect E2E Tests 3" - uses: ./.github/workflows/e2e-run.yml - if: ${{ github.repository_owner == 'AutoMQ' }} - with: - suite-id: "connect3" - test-yaml: "tests/suites/connect_test_suite3.yml" - runner: "e2e" - streams_e2e: - name: "Run streams E2E Tests" - uses: ./.github/workflows/e2e-run.yml - if: ${{ github.repository_owner == 'AutoMQ' }} - with: - suite-id: "streams" - test-path: "tests/kafkatest/tests/streams" - runner: "e2e" - e2e_summary: - name: "E2E Tests Summary" - runs-on: "e2e" - if: ${{ always() && github.repository_owner == 'AutoMQ' }} - needs: [ benchmarks_e2e, connect_e2e_1, connect_e2e_2, connect_e2e_3, streams_e2e ] - steps: - - name: Report results - run: python3 tests/report_e2e_results.py - env: - CURRENT_REPO: ${{ github.repository }} - RUN_ID: ${{ github.run_id }} - WEB_HOOK_URL: ${{ secrets.E2E_REPORT_WEB_HOOK_URL }} - DATA_MAP: "{\"benchmarks_e2e\": ${{ toJSON(needs.benchmarks_e2e.outputs) }}, \"connect_e2e_1\": ${{ toJSON(needs.connect_e2e_1.outputs) }}, \"connect_e2e_2\": ${{ toJSON(needs.connect_e2e_2.outputs) }}, \"connect_e2e_3\": ${{ toJSON(needs.connect_e2e_3.outputs) }}, \"streams_e2e\": ${{ toJSON(needs.streams_e2e.outputs) }}}" - REPORT_TITLE_PREFIX: "Extra" diff --git a/.github/workflows/nightly-main-e2e.yml b/.github/workflows/nightly-main-e2e.yml index 4bc7fe296f..d0855b232a 100644 --- a/.github/workflows/nightly-main-e2e.yml +++ b/.github/workflows/nightly-main-e2e.yml @@ -1,8 +1,8 @@ -name: Nightly Main E2E tests +name: Nightly E2E tests on: workflow_dispatch: schedule: - - cron: '0 16 * * *' + - cron: '0 16 1,7,14,21,28 * *' jobs: main_e2e_1: @@ -45,11 +45,51 @@ jobs: suite-id: "main5" test-path: "tests/kafkatest/automq" runner: "e2e" + benchmarks_e2e: + name: "Run benchmarks E2E Tests" + uses: ./.github/workflows/e2e-run.yml + if: ${{ github.repository_owner == 'AutoMQ' }} + with: + suite-id: "benchmarks" + test-path: "tests/kafkatest/benchmarks" + runner: "e2e" + connect_e2e_1: + name: "Run connect E2E Tests 1" + uses: ./.github/workflows/e2e-run.yml + if: ${{ github.repository_owner == 'AutoMQ' }} + with: + suite-id: "connect1" + test-yaml: "tests/suites/connect_test_suite1.yml" + runner: "e2e" + connect_e2e_2: + name: "Run connect E2E Tests 2" + uses: ./.github/workflows/e2e-run.yml + if: ${{ github.repository_owner == 'AutoMQ' }} + with: + suite-id: "connect2" + test-yaml: "tests/suites/connect_test_suite2.yml" + runner: "e2e" + connect_e2e_3: + name: "Run connect E2E Tests 3" + uses: ./.github/workflows/e2e-run.yml + if: ${{ github.repository_owner == 'AutoMQ' }} + with: + suite-id: "connect3" + test-yaml: "tests/suites/connect_test_suite3.yml" + runner: "e2e" + streams_e2e: + name: "Run streams E2E Tests" + uses: ./.github/workflows/e2e-run.yml + if: ${{ github.repository_owner == 'AutoMQ' }} + with: + suite-id: "streams" + test-path: "tests/kafkatest/tests/streams" + runner: "e2e" e2e_summary: runs-on: "e2e" name: "E2E Tests Summary" if: ${{ always() && github.repository_owner == 'AutoMQ' }} - needs: [ main_e2e_1, main_e2e_2, main_e2e_3, main_e2e_4, main_e2e_5 ] + needs: [ main_e2e_1, main_e2e_2, main_e2e_3, main_e2e_4, main_e2e_5, benchmarks_e2e, connect_e2e_1, connect_e2e_2, connect_e2e_3, streams_e2e ] steps: - name: Report results run: python3 tests/report_e2e_results.py @@ -57,5 +97,5 @@ jobs: CURRENT_REPO: ${{ github.repository }} RUN_ID: ${{ github.run_id }} WEB_HOOK_URL: ${{ secrets.E2E_REPORT_WEB_HOOK_URL }} - DATA_MAP: "{\"main_e2e_1\": ${{ toJSON(needs.main_e2e_1.outputs) }}, \"main_e2e_2\": ${{ toJSON(needs.main_e2e_2.outputs) }}, \"main_e2e_3\": ${{ toJSON(needs.main_e2e_3.outputs) }}, \"main_e2e_4\": ${{ toJSON(needs.main_e2e_4.outputs) }}, \"main_e2e_5\": ${{ toJSON(needs.main_e2e_5.outputs) }}}" + DATA_MAP: "{\"main_e2e_1\": ${{ toJSON(needs.main_e2e_1.outputs) }}, \"main_e2e_2\": ${{ toJSON(needs.main_e2e_2.outputs) }}, \"main_e2e_3\": ${{ toJSON(needs.main_e2e_3.outputs) }}, \"main_e2e_4\": ${{ toJSON(needs.main_e2e_4.outputs) }}, \"main_e2e_5\": ${{ toJSON(needs.main_e2e_5.outputs) }}, \"benchmarks_e2e\": ${{ toJSON(needs.benchmarks_e2e.outputs) }}, \"connect_e2e_1\": ${{ toJSON(needs.connect_e2e_1.outputs) }}, \"connect_e2e_2\": ${{ toJSON(needs.connect_e2e_2.outputs) }}, \"connect_e2e_3\": ${{ toJSON(needs.connect_e2e_3.outputs) }}, \"streams_e2e\": ${{ toJSON(needs.streams_e2e.outputs) }}}" REPORT_TITLE_PREFIX: "Main" diff --git a/.github/workflows/publish_maven_package.yml b/.github/workflows/publish_maven_package.yml new file mode 100644 index 0000000000..bbec2e5b5f --- /dev/null +++ b/.github/workflows/publish_maven_package.yml @@ -0,0 +1,59 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: Publish Maven Package + +on: + workflow_dispatch: + inputs: + version: + description: 'Version to publish' + required: true + push: + tags: + - '[0-9]+.[0-9]+.[0-9]+' + - '[0-9]+.[0-9]+.[0-9]+-rc[0-9]+' + +env: + VERSION: ${{ github.event.inputs.version || github.ref_name }} + +jobs: + publish: + name: "Publish to Github Packages" + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ ubuntu-22.04 ] + jdk: [ 17 ] + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Gradle wrapper validation + uses: gradle/actions/wrapper-validation@v3 + - name: Set up JDK ${{ matrix.jdk }} + uses: actions/setup-java@v3 + with: + java-version: ${{ matrix.jdk }} + distribution: "zulu" + - name: Setup Gradle + uses: gradle/actions/setup-gradle@v4 + with: + gradle-version: '8.10' + - name: Publish + run: | + gradle publish -PmavenUrl='https://maven.pkg.github.com/AutoMQ/automq' \ + -PmavenUsername=${{ env.GITHUB_ACTOR }} -PmavenPassword=${{ secrets.GITHUB_TOKEN }} \ + -PskipSigning=true \ + -Pgroup=com.automq.automq -Pversion=${{ env.VERSION }} diff --git a/.github/workflows/s3stream-e2e.yml b/.github/workflows/s3stream-e2e.yml deleted file mode 100644 index 260332264e..0000000000 --- a/.github/workflows/s3stream-e2e.yml +++ /dev/null @@ -1,50 +0,0 @@ -name: E2E-TEST for AutoMQ S3Stream -on: - push: - branches: [ "main" ] - workflow_dispatch: - -jobs: - test: - runs-on: ubuntu-latest - strategy: - matrix: - wal-type: [ "file", "s3" ] - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-java@v4 - with: - distribution: 'temurin' - java-version: '17' - cache: "maven" - - name: Build s3stream - working-directory: ./s3stream - run: mvn clean install -U -DskipTests - - uses: actions/checkout@v3 - with: - repository: AutoMQ/s3stream-e2e - ref: main - path: s3stream-e2e - token: ${{ secrets.PAT_E2E }} - - name: Change s3stream version - run: | - export S3STREAM_VERSION=$(cat s3stream/target/maven-archiver/pom.properties | grep version | awk -F '=' '{print $2}') - echo "change s3stream version of e2e test to $S3STREAM_VERSION" - sed -i "s/.*<\/s3stream.version>/$S3STREAM_VERSION<\/s3stream.version>/g" s3stream-e2e/pom.xml - export LINE_START=$(awk '//{print NR}' s3stream-e2e/pom.xml) - export LINE_END=$(awk '/<\/repositories>/{print NR}' s3stream-e2e/pom.xml) - sed -i "${LINE_START},${LINE_END}d" s3stream-e2e/pom.xml - cat s3stream-e2e/pom.xml - - name: Run tests - working-directory: ./s3stream-e2e - run: mvn test -pl integration - env: - WAL_TYPE: ${{ matrix.wal-type }} - - name: Publish Test Report - uses: mikepenz/action-junit-report@v3 - if: success() || failure() # always run even if the previous step fails - with: - report_paths: '**/surefire-reports/TEST-*.xml' - annotate_only: true - include_passed: true - detailed_summary: true diff --git a/.github/workflows/spark-iceberg-docker-build.yml b/.github/workflows/spark-iceberg-docker-build.yml new file mode 100644 index 0000000000..0533bbc1d3 --- /dev/null +++ b/.github/workflows/spark-iceberg-docker-build.yml @@ -0,0 +1,31 @@ +name: Spark Iceberg image + +on: + workflow_dispatch: + +jobs: + docker: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_READ_WRITE_TOKEN }} + + - name: Build and Push + uses: docker/build-push-action@v6 + with: + context: docker/table_topic/spark_iceberg/ + platforms: linux/amd64,linux/arm64 + push: true + tags: automqinc/spark-iceberg:latest diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index ec01e252ee..3da8cae57d 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -43,13 +43,14 @@ jobs: debug-only: ${{ inputs.dryRun || false }} operations-per-run: ${{ inputs.operationsPerRun || 100 }} days-before-stale: 90 - days-before-close: -1 + days-before-close: 30 + days-before-issue-stale: -1 + days-before-issue-close: -1 stale-pr-label: 'stale' stale-pr-message: > This PR is being marked as stale since it has not had any activity in 90 days. If you - would like to keep this PR alive, please ask a committer for review. If the PR has + would like to keep this PR alive, please ask a committer for review. If the PR has merge conflicts, please update it with the latest from trunk (or appropriate release branch)

If this PR is no longer valid or desired, please feel free to close it. If no activity occurs in the next 30 days, it will be automatically closed. - diff --git a/.gitignore b/.gitignore index 015df8ead8..1e3f406d25 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,7 @@ TAGS Vagrantfile.local /logs .DS_Store +local.properties config/server-* config/zookeeper-* @@ -62,3 +63,7 @@ storage/kafka-tiered-storage/ docker/test/report_*.html kafka.Kafka __pycache__ + +# Ignore bin folder generated by the build, but exclude the one in the root +bin/ +!/bin/ diff --git a/.idea/copyright/BSL.xml b/.idea/copyright/BSL.xml deleted file mode 100644 index e1358f8481..0000000000 --- a/.idea/copyright/BSL.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/copyright/profiles_settings.xml b/.idea/copyright/profiles_settings.xml deleted file mode 100644 index 2fd6b1128b..0000000000 --- a/.idea/copyright/profiles_settings.xml +++ /dev/null @@ -1,7 +0,0 @@ - - - - - - - \ No newline at end of file diff --git a/CONTRIBUTING_GUIDE.md b/CONTRIBUTING_GUIDE.md index 0988e5a57f..027cf0f1aa 100644 --- a/CONTRIBUTING_GUIDE.md +++ b/CONTRIBUTING_GUIDE.md @@ -2,18 +2,24 @@ Thank you for your interest in contributing! We love community contributions. Read on to learn how to contribute to AutoMQ. -We appreciate first time contributors and we are happy to assist you in getting started. In case of questions, just +We appreciate first-time contributors, and we are happy to assist you in getting started. In case of questions, just reach out to us via [Wechat Group](https://www.automq.com/img/----------------------------1.png) or [Slack](https://join.slack.com/t/automq/shared_invite/zt-29h17vye9-thf31ebIVL9oXuRdACnOIA)! -Before getting started, please review AutoMQ's Code of Conduct. Everyone interacting in Slack or Wechat +Before getting started, please review AutoMQ's Code of Conduct. Everyone interacting in Slack or WeChat follow [Code of Conduct](CODE_OF_CONDUCT.md). ## Code Contributions -Most of the issues open for contributions are tagged with 'good first issue.' To claim one, simply reply with 'pick up' in the issue and the AutoMQ maintainers will assign the issue to you. If you have any questions about the 'good first issue' please feel free to ask. We will do our best to clarify any doubts you may have. -Start with -this [tagged good first issue](https://github.com/AutoMQ/automq-for-kafka/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) +### Finding or Reporting Issues + +- **Find an existing issue:** Look through the [existing issues](https://github.com/AutoMQ/automq/issues). Issues open for contributions are often tagged with `good first issue`. To claim an issue, simply reply with '/assign', and the GitHub bot will assign it to you. Start with + this [tagged good first issue](https://github.com/AutoMQ/automq-for-kafka/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22). +- **Report a new issue:** If you've found a bug or have a feature request, please [create a new issue](https://github.com/AutoMQ/automq/issues/new/choose). Select the appropriate template (Bug Report or Feature Request) and fill out the form provided. + +If you have any questions about an issue, please feel free to ask in the issue comments. We will do our best to clarify any doubts you may have. + +### Submitting Pull Requests The usual workflow of code contribution is: @@ -25,24 +31,24 @@ The usual workflow of code contribution is: 5. Push your local branch to your fork. 6. Submit a Pull Request so that we can review your changes. 7. [Link an existing Issue](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue) - that does not include the `needs triage` label to your Pull Request. A pull request without a linked issue will be + (created via the steps above or an existing one you claimed) that does not include the `needs triage` label to your Pull Request. A pull request without a linked issue will be closed, otherwise. 8. Write a PR title and description that follows the [Pull Request Template](PULL_REQUEST_TEMPLATE.md). 9. An AutoMQ maintainer will trigger the CI tests for you and review the code. -10. Review and respond to feedback and questions by AutoMQ maintainers. +10. Review and respond to feedback and questions from AutoMQ maintainers. 11. Merge the contribution. Pull Request reviews are done on a regular basis. -> [!NOTE] +> [!NOTE] > Please make sure you respond to our feedback/questions and sign our CLA. > -> Pull Requests without updates will be closed due inactivity. +> Pull Requests without updates will be closed due to inactivity. ## Requirement | Requirement | Version | -|------------------------|------------| +| ---------------------- | ---------- | | Compiling requirements | JDK 17 | | Compiling requirements | Scala 2.13 | | Running requirements | JDK 17 | @@ -58,17 +64,21 @@ Building AutoMQ is the same as Apache Kafka. Kafka uses Gradle as its project ma It is not recommended to manually install Gradle. The gradlew script in the root directory will automatically download Gradle for you, and the version is also specified by the gradlew script. ### Build + ``` ./gradlew jar -x test ``` ### Prepare S3 service -Refer to this [documentation](https://docs.localstack.cloud/getting-started/installation/) to install `localstack` to mock a local s3 service or use AWS S3 service directly. + +Refer to this [documentation](https://docs.localstack.cloud/getting-started/installation/) to install `localstack` to mock a local S3 service or use AWS S3 service directly. If you are using localstack then create a bucket with the following command: + ``` aws s3api create-bucket --bucket ko3 --endpoint=http://127.0.0.1:4566 ``` + ### Modify Configuration Modify the `config/kraft/server.properties` file. The following settings need to be changed: @@ -83,28 +93,34 @@ s3.region=us-east-1 # The bucket of S3 service to store data s3.bucket=ko3 ``` + > Tips: If you're using localstack, make sure to set the s3.endpoint to http://127.0.0.1:4566, not localhost. Set the region to us-east-1. The bucket should match the one created earlier. ### Format + Generated Cluster UUID: + ``` KAFKA_CLUSTER_ID="$(bin/kafka-storage.sh random-uuid)" ``` + Format Metadata Catalog: + ``` bin/kafka-storage.sh format -t $KAFKA_CLUSTER_ID -c config/kraft/server.properties ``` + ### IDE Start Configuration -| Item | Value | -|------------------------|------------| -| Main | core/src/main/scala/kafka/Kafka.scala | -| ClassPath | -cp kafka.core.main | -| VM Options | -Xmx1 -Xms1G -server -XX:+UseZGC -XX:MaxDirectMemorySize=2G -Dkafka.logs.dir=logs/ -Dlog4j.configuration=file:config/log4j.properties -Dio.netty.leakDetection.level=paranoid | -| CLI Arguments | config/kraft/server.properties| -| Environment | KAFKA_S3_ACCESS_KEY=test;KAFKA_S3_SECRET_KEY=test | -> tips: If you are using localstack, just use any value of access key and secret key. If you are using real S3 service, set `KAFKA_S3_ACCESS_KEY` and `KAFKA_S3_SECRET_KEY` to the real access key and secret key that have read/write permission of S3 service. +| Item | Value | +| ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| Main | core/src/main/scala/kafka/Kafka.scala | +| ClassPath | -cp kafka.core.main | +| VM Options | -Xmx1G -Xms1G -server -XX:+UseZGC -XX:MaxDirectMemorySize=2G -Dkafka.logs.dir=logs/ -Dlog4j.configuration=file:config/log4j.properties -Dio.netty.leakDetection.level=paranoid | +| CLI Arguments | config/kraft/server.properties | +| Environment | KAFKA_S3_ACCESS_KEY=test;KAFKA_S3_SECRET_KEY=test | +> tips: If you are using localstack, just use any value of access key and secret key. If you are using real S3 service, set `KAFKA_S3_ACCESS_KEY` and `KAFKA_S3_SECRET_KEY` to the real access key and secret key that have read/write permission of S3 service. ## Documentation diff --git a/LICENSE b/LICENSE index ffbfecdab9..d645695673 100644 --- a/LICENSE +++ b/LICENSE @@ -1,29 +1,202 @@ -Copyright (c) 2023-2024 AutoMQ HK Limited. - -this software are licensed as follows: - -1. Apache Kafka Source and Dependency Licensing: - All code in this repository that is forked from Apache Kafka and its - dependencies will continue to be licensed under the original Apache Kafka - open source license. For detailed licensing information regarding Apache - Kafka and its dependencies, please refer to the files under the "/licenses/" - folder in this repository. - -2. S3Stream Component Licensing: - The S3Stream component added to this project (specifically referring to all - files under the "/S3Stream/" directory) is licensed under a revised Business - Source License (BSL) by AutoMQ HK Limited, with the specific terms available - in the /LICENSE.S3Stream file in this repository. Any dependencies used by - the S3Stream component are subject to their respective open source licenses. - -3. File-Level License Precedence: - For each file in this repository, if the license is explicitly specified in - the header of the file, the license stated in the file header shall prevail. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/LICENSE.S3Stream b/LICENSE.S3Stream deleted file mode 100644 index 30f7e98db3..0000000000 --- a/LICENSE.S3Stream +++ /dev/null @@ -1,96 +0,0 @@ -License text copyright © 2023 MariaDB plc, All Rights Reserved. -"Business Source License" is a trademark of MariaDB plc. - - -Parameters - -Licensor: AutoMQ HK Limited. -Licensed Work: AutoMQ Version 1.1.2 or later. The Licensed Work is (c) 2024 - AutoMQ HK Limited. -Additional Use Grant: You may make production use of the Licensed Work, provided - Your use does not include offering the Licensed Work to third - parties on a hosted or embedded basis in order to compete with - AutoMQ's paid version(s) of the Licensed Work. For purposes - of this license: - - A "competitive offering" is a Product that is offered to third - parties on a paid basis, including through paid support - arrangements, that significantly overlaps with the capabilities - of AutoMQ's paid version(s) of the Licensed Work. If Your - Product is not a competitive offering when You first make it - generally available, it will not become a competitive offering - later due to AutoMQ releasing a new version of the Licensed - Work with additional capabilities. In addition, Products that - are not provided on a paid basis are not competitive. - - "Product" means software that is offered to end users to manage - in their own environments or offered as a service on a hosted - basis. - - "Embedded" means including the source code or executable code - from the Licensed Work in a competitive offering. "Embedded" - also means packaging the competitive offering in such a way - that the Licensed Work must be accessed or downloaded for the - competitive offering to operate. - - Hosting or using the Licensed Work(s) for internal purposes - within an organization is not considered a competitive - offering. AutoMQ considers your organization to include all - of your affiliates under common control. - - For binding interpretive guidance on using AutoMQ products - under the Business Source License, please visit our FAQ. - (https://www.automq.com/license-faq) -Change Date: Change date is four years from release date. - Please see https://github.com/AutoMQ/automq/releases for exact dates -Change License: Apache License, Version 2.0 - URL: https://www.apache.org/licenses/LICENSE-2.0 - - -For information about alternative licensing arrangements for the Licensed Work, -please contact licensing@automq.com. - -Notice - -Business Source License 1.1 - -Terms - -The Licensor hereby grants you the right to copy, modify, create derivative -works, redistribute, and make non-production use of the Licensed Work. The -Licensor may make an Additional Use Grant, above, permitting limited production use. - -Effective on the Change Date, or the fourth anniversary of the first publicly -available distribution of a specific version of the Licensed Work under this -License, whichever comes first, the Licensor hereby grants you rights under -the terms of the Change License, and the rights granted in the paragraph -above terminate. - -If your use of the Licensed Work does not comply with the requirements -currently in effect as described in this License, you must purchase a -commercial license from the Licensor, its affiliated entities, or authorized -resellers, or you must refrain from using the Licensed Work. - -All copies of the original and modified Licensed Work, and derivative works -of the Licensed Work, are subject to this License. This License applies -separately for each version of the Licensed Work and the Change Date may vary -for each version of the Licensed Work released by Licensor. - -You must conspicuously display this License on each original or modified copy -of the Licensed Work. If you receive the Licensed Work in original or -modified form from a third party, the terms and conditions set forth in this -License apply to your use of that work. - -Any use of the Licensed Work in violation of this License will automatically -terminate your rights under this License for the current and all other -versions of the Licensed Work. - -This License does not grant you any right in any trademark or logo of -Licensor or its affiliates (provided that you may use a trademark or logo of -Licensor as expressly required by this License). - -TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON -AN "AS IS" BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS, -EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF -MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND -TITLE. diff --git a/NOTICE b/NOTICE index a9bf726a67..483fd2d3d1 100644 --- a/NOTICE +++ b/NOTICE @@ -1,5 +1,5 @@ AutoMQ NOTICE -Copyright 2023-2024, AutoMQ HK Limited. +Copyright 2023-2025, AutoMQ HK Limited. --------------------------- Apache Kafka NOTICE diff --git a/NOTICE-binary b/NOTICE-binary index 988a33e7fe..3ba6e76578 100644 --- a/NOTICE-binary +++ b/NOTICE-binary @@ -1,5 +1,5 @@ AutoMQ Binary NOTICE -Copyright 2023-2024, AutoMQ HK Limited. +Copyright 2023-2025, AutoMQ HK Limited. --------------------------- Apache Kafka Binary NOTICE diff --git a/README.md b/README.md index b37945a423..bca79f039e 100644 --- a/README.md +++ b/README.md @@ -1,233 +1,177 @@ -# AutoMQ: A cloud-first alternative of Kafka by decoupling durability to S3 and EBS - +# A Diskless Kafka® on S3, Offering 10x Cost Savings and Scaling in Seconds.

- 🔥  Free trial of AutoMQ Business Edition    📑  Documentation    - 📃  AutoMQ Introduction + >Free trial of AutoMQ on AWS   

- [![Linkedin Badge](https://img.shields.io/badge/-LinkedIn-blue?style=flat-square&logo=Linkedin&logoColor=white&link=https://www.linkedin.com/company/automq)](https://www.linkedin.com/company/automq) -[![Twitter URL](https://img.shields.io/twitter/follow/AutoMQ)](https://twitter.com/intent/follow?screen_name=AutoMQ_Lab) -[![](https://img.shields.io/badge/-%20Wechat%20-red?style=social&logo=discourse)](docs/images/automq-wechat.png) -[![](https://badgen.net/badge/Slack/Join%20AutoMQ/0abd59?icon=slack)](https://join.slack.com/t/automq/shared_invite/zt-29h17vye9-thf31ebIVL9oXuRdACnOIA) -[![](https://img.shields.io/badge/AutoMQ%20vs.%20Kafka(Cost)-yellow)](https://www.automq.com/blog/automq-vs-apache-kafka-a-real-aws-cloud-bill-comparison) -[![](https://img.shields.io/badge/AutoMQ%20vs.%20Kafka(Performance)-orange)](https://docs.automq.com/docs/automq-opensource/IJLQwnVROiS5cUkXfF0cuHnWnNd) +[![](https://badgen.net/badge/Slack/Join%20AutoMQ/0abd59?icon=slack)](https://go.automq.com/slack) +[![](https://img.shields.io/badge/AutoMQ%20vs.%20Kafka(Cost)-yellow)](https://www.automq.com/blog/automq-vs-apache-kafka-a-real-aws-cloud-bill-comparison?utm_source=github_automq) +[![](https://img.shields.io/badge/AutoMQ%20vs.%20Kafka(Performance)-orange)](https://www.automq.com/docs/automq/benchmarks/automq-vs-apache-kafka-benchmarks-and-cost?utm_source=github_automq) +[![Gurubase](https://img.shields.io/badge/Gurubase-Ask%20AutoMQ%20Guru-006BFF)](https://gurubase.io/g/automq) +[![DeepWiki](https://img.shields.io/badge/DeepWiki-AutoMQ%2Fautomq-blue.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACwAAAAyCAYAAAAnWDnqAAAAAXNSR0IArs4c6QAAA05JREFUaEPtmUtyEzEQhtWTQyQLHNak2AB7ZnyXZMEjXMGeK/AIi+QuHrMnbChYY7MIh8g01fJoopFb0uhhEqqcbWTp06/uv1saEDv4O3n3dV60RfP947Mm9/SQc0ICFQgzfc4CYZoTPAswgSJCCUJUnAAoRHOAUOcATwbmVLWdGoH//PB8mnKqScAhsD0kYP3j/Yt5LPQe2KvcXmGvRHcDnpxfL2zOYJ1mFwrryWTz0advv1Ut4CJgf5uhDuDj5eUcAUoahrdY/56ebRWeraTjMt/00Sh3UDtjgHtQNHwcRGOC98BJEAEymycmYcWwOprTgcB6VZ5JK5TAJ+fXGLBm3FDAmn6oPPjR4rKCAoJCal2eAiQp2x0vxTPB3ALO2CRkwmDy5WohzBDwSEFKRwPbknEggCPB/imwrycgxX2NzoMCHhPkDwqYMr9tRcP5qNrMZHkVnOjRMWwLCcr8ohBVb1OMjxLwGCvjTikrsBOiA6fNyCrm8V1rP93iVPpwaE+gO0SsWmPiXB+jikdf6SizrT5qKasx5j8ABbHpFTx+vFXp9EnYQmLx02h1QTTrl6eDqxLnGjporxl3NL3agEvXdT0WmEost648sQOYAeJS9Q7bfUVoMGnjo4AZdUMQku50McDcMWcBPvr0SzbTAFDfvJqwLzgxwATnCgnp4wDl6Aa+Ax283gghmj+vj7feE2KBBRMW3FzOpLOADl0Isb5587h/U4gGvkt5v60Z1VLG8BhYjbzRwyQZemwAd6cCR5/XFWLYZRIMpX39AR0tjaGGiGzLVyhse5C9RKC6ai42ppWPKiBagOvaYk8lO7DajerabOZP46Lby5wKjw1HCRx7p9sVMOWGzb/vA1hwiWc6jm3MvQDTogQkiqIhJV0nBQBTU+3okKCFDy9WwferkHjtxib7t3xIUQtHxnIwtx4mpg26/HfwVNVDb4oI9RHmx5WGelRVlrtiw43zboCLaxv46AZeB3IlTkwouebTr1y2NjSpHz68WNFjHvupy3q8TFn3Hos2IAk4Ju5dCo8B3wP7VPr/FGaKiG+T+v+TQqIrOqMTL1VdWV1DdmcbO8KXBz6esmYWYKPwDL5b5FA1a0hwapHiom0r/cKaoqr+27/XcrS5UwSMbQAAAABJRU5ErkJggg==)](https://deepwiki.com/AutoMQ/automq) AutoMQ%2Fautomq | Trendshift +
---- +
-![](https://img.shields.io/badge/AWS-%E2%9C%85-lightgray?logo=amazonaws) -![](https://img.shields.io/badge/Google-%F0%9F%9A%A7-lightyellow?logo=googlecloud) -![](https://img.shields.io/badge/Azure-%F0%9F%9A%A7-lightyellow?logo=microsoftazure) -![](https://img.shields.io/badge/Aliyun-%E2%9C%85-lightgray?logo=alibabacloud) -![](https://img.shields.io/badge/Huawei-%E2%9C%85-lightgray?logo=huawei) -![](https://img.shields.io/badge/Baidu-%E2%9C%85-lightgray?logo=baidu) -![](https://img.shields.io/badge/Tencent-%E2%9C%85-lightgray?logo=tencentqq) +automq-solgan + + + Grab + + + + Avia + + + Tencent + + + Honda + + + Trip + + + LG + + + JD + + + + Geely + + + Poizon + + + Bitkub + + + PalmPay + + + RedNote + + + XPENG + + + OPPO + + + BambuLab +
-## 📺 Youtube Video Introduction -Watch this video to learn what is AutoMQ. ⬇️ ⬇️ ⬇️ - -[![What is AutoMQ?](https://img.youtube.com/vi/3JQrclZlie4/0.jpg)](https://www.youtube.com/watch?v=3JQrclZlie4) - - - -## 🍵 AutoMQ vs Other Streaming Platforms - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
FeatureAutoMQApache KafkaConfluentApache PulsarRedpandaWarpstream
Apache Kafka Compatibility[1]Native KafkaNon-KafkaKafka Protocol
Source Code AvailabilityYesYesNoYesYesNo
Stateless BrokerYesNoNoYesNoYes
Publisher Latency(P99)Single-digit ms latency> 620ms
Continuous Self-BalancingYesNoYesYesYesYes
Scale in/outIn secondsIn hours/daysIn hoursIn hours
(scale-in);
In seconds
(scale-out)		In hours
In seconds (Enterprise Only)		In seconds
Spot Instance SupportYesNoNoNoNoYes
Partition ReassignmentIn secondsIn hours/daysIn hoursIn secondsIn hours
In seconds (Enterprise Only)		In seconds
ComponentBrokerBroker
Zookeeper
(Non-KRaft)		Broker
Zookeeper
Bookkeeper
Proxy(Optional)		BrokerAgent
MetadataServer
DurabilityGuaranteed by S3/EBS[2]Guaranteed by ISR Guaranteed by BookkeeperGuaranteed by RaftGuaranteed by S3
Inter-AZ Networking FeesNoYesNo
- - -> [1] Apache Kafka Compatibility's definition is coming from this [blog](https://www.kai-waehner.de/blog/2021/05/09/kafka-api-de-facto-standard-event-streaming-like-amazon-s3-object-storage/). - -> [2] EBS Durability: On Azure, GCP, and Alibaba Cloud, Regional EBS replicas span multiple AZs. On AWS, ensure durability by double writing to EBS and S3 Express One Zone in different AZs. + +- [Grab: Driving Efficiency with AutoMQ in DataStreaming Platform](https://www.youtube.com/watch?v=IB8sh639Rsg) +- [JD.com x AutoMQ x CubeFS: A Cost-Effective Journey](https://www.automq.com/blog/jdcom-automq-cubefs-trillion-scale-kafka-messaging?utm_source=github_automq) +- [Palmpay Uses AutoMQ to Replace Kafka, Optimizing Costs by 50%+](https://www.automq.com/blog/palmpay-uses-automq-to-replace-kafka?utm_source=github_automq) +- [AutoMQ help Geely Auto(Fortune Global 500) solve the pain points of Kafka elasticity in the V2X scenario](https://www.automq.com/blog/automq-help-geely-auto-solve-the-pain-points-of-kafka-elasticity-in-the-v2x-scenario?utm_source=github_automq) +- [How Asia’s Quora Zhihu uses AutoMQ to reduce Kafka cost and maintenance complexity](https://www.automq.com/blog/how-asias-quora-zhihu-use-automq-to-reduce-kafka-cost-and-maintenance-complexity?utm_source=github_automq) +- [XPENG Motors Reduces Costs by 50%+ by Replacing Kafka with AutoMQ](https://www.automq.com/blog/xpeng-motors-reduces-costs-by-50-by-replacing-kafka-with-automq?utm_source=github_automq) +- [Asia's GOAT, Poizon uses AutoMQ Kafka to build observability platform for massive data(30 GB/s)](https://www.automq.com/blog/asiax27s-goat-poizon-uses-automq-kafka-to-build-a-new-generation-observability-platform-for-massive-data?utm_source=github_automq) +- [AutoMQ Helps CaoCao Mobility Address Kafka Scalability During Holidays](https://www.automq.com/blog/automq-helps-caocao-mobility-address-kafka-scalability-issues-during-mid-autumn-and-national-day?utm_source=github_automq) + +## ⛄ Get started with AutoMQ + +> [!Tip] +> Deploying a production-ready AutoMQ cluster is challenging. This Quick Start is only for evaluating AutoMQ features and is not suitable for production use. For production deployment best practices, please [contact](https://www.automq.com/contact) our community for support. + +The `docker/docker-compose.yaml` file provides a simple single-node setup for quick evaluation and development: +```shell +curl -O https://raw.githubusercontent.com/AutoMQ/automq/refs/tags/1.5.5/docker/docker-compose.yaml && docker compose -f docker-compose.yaml up -d +``` +This setup features a single AutoMQ node serving as both controller and broker, alongside MinIO for S3 storage. All services operate within a Docker bridge network called `automq_net`, allowing you to start a Kafka producer in this network to test AutoMQ: +```shell +docker run --network automq_net automqinc/automq:latest /bin/bash -c \ +"/opt/automq/kafka/bin/kafka-producer-perf-test.sh --topic test-topic --num-records=1024000 --throughput 5120 --record-size 1024 \ +--producer-props bootstrap.servers=server1:9092 linger.ms=100 batch.size=524288 buffer.memory=134217728 max.request.size=67108864" +``` +After testing, you can destroy the setup with: +```shell +docker compose -f docker-compose.yaml down +``` +The `docker/docker-compose-cluster.yaml` file offers a more complex setup with three AutoMQ nodes, ideal for testing AutoMQ's cluster features, and can be run in the same way. + +There are more deployment options available: +- [Deploy Multi-Nodes Test Cluster on Docker](https://www.automq.com/docs/automq/getting-started/deploy-multi-nodes-test-cluster-on-docker?utm_source=github_automq) +- [Deploy on Linux with 5 Nodes](https://www.automq.com/docs/automq/deployment/deploy-multi-nodes-cluster-on-linux?utm_source=github_automq) +- [Deploy on Kubernetes](https://www.automq.com/docs/automq/deployment/deploy-multi-nodes-cluster-on-kubernetes?utm_source=github_automq) +- [Try AutoMQ on AWS Marketplace (Two Weeks Free Trial)](https://docs.automq.com/automq-cloud/getting-started/install-byoc-environment/aws/install-env-from-marketplace?utm_source=github_automq) +- [Try AutoMQ on Alibaba Cloud Marketplace (Two Weeks Free Trial)](https://market.aliyun.com/products/55530001/cmgj00065841.html) + +## 🗞️ Newest Feature - Table Topic +Table Topic is a new feature in AutoMQ that combines stream and table functionalities to unify streaming and data analysis. Currently, it supports Apache Iceberg and integrates with catalog services such as AWS Glue, HMS, and the Rest catalog. Additionally, it natively supports S3 tables, a new AWS product announced at the 2024 re:Invent. [Learn more](https://www.automq.com/blog/automq-table-topic-seamless-integration-with-s3-tables-and-iceberg?utm_source=github_automq). + +![image](https://github.com/user-attachments/assets/6b2a514a-cc3e-442e-84f6-d953206865e0) ## 🔶 Why AutoMQ +AutoMQ is a stateless Kafka alternative that runs on S3 or any S3-compatible storage, such as MinIO. It is designed to address two major issues of Apache Kafka. First, Kafka clusters are difficult to scale out or in due to the stateful nature of its brokers. Data movement is required, and even reassigning partitions between brokers is a complex process. Second, hosting Kafka in the cloud can be prohibitively expensive. You face high costs for EBS storage, cross-AZ traffic, and significant over-provisioning due to Kafka's limited scalability. -- **Cost effective**: The first true cloud-native streaming storage system, designed for optimal cost and efficiency on the cloud. Refer to [this report](https://docs.automq.com/docs/automq-opensource/EV6mwoC95ihwRckMsUKcppnqnJb) to see how we cut Apache Kafka billing by 90% on the cloud. -- **High Reliability**: Leverage cloud-shared storage services(EBS and S3) to achieve zero RPO, RTO in seconds and 99.999999999% durability. +Here are some key highlights of AutoMQ that make it an ideal choice to replace your Apache Kafka cluster, whether in the cloud or on-premise, as long as you have S3-compatible storage: +- **Cost effective**: The first true cloud-native streaming storage system, designed for optimal cost and efficiency on the cloud. Refer to [this report](https://www.automq.com/docs/automq/benchmarks/cost-effective-automq-vs-apache-kafka?utm_source=github_automq) to see how we cut Apache Kafka billing by 90% on the cloud. +- **High Reliability**: Leverage object storage service to achieve zero RPO, RTO in seconds and 99.999999999% durability. +- **Zero Cross-AZ Traffic**: By using cloud object storage as the priority storage solution, AutoMQ eliminates cross-AZ traffic costs on AWS and GCP. In traditional Kafka setups, over 80% of costs arise from cross-AZ traffic, including producer, consumer, and replication sides. - **Serverless**: - - Auto Scaling: Monitor cluster metrics and automatically scale in/out to align with your workload, enabling a pay-as-you-go model. - - Scaling in seconds: The computing layer (broker) is stateless and can scale in/out within seconds, making AutoMQ a truly serverless solution. - - Infinite scalable: Utilize cloud object storage as the primary storage solution, eliminating concerns about storage capacity. + - Auto Scaling: Monitor cluster metrics and automatically scale in/out to align with your workload, enabling a pay-as-you-go model. + - Scaling in seconds: The computing layer (broker) is stateless and can scale in/out within seconds, making AutoMQ a truly serverless solution. + - Infinite scalable: Utilize cloud object storage as the primary storage solution, eliminating concerns about storage capacity. - **Manage-less**: The built-in auto-balancer component automatically schedules partitions and network traffic between brokers, eliminating manual partition reassignment. - **High performance**: - - Low latency: Accelerate writing with high-performance EBS as WAL, achieving single-digit millisecond latency. - - High throughput: Leverage pre-fetching, batch processing, and parallel technologies to maximize the capabilities of cloud object storage. - > Refer to the [AutoMQ Performance White Paper](https://docs.automq.com/docs/automq-opensource/IJLQwnVROiS5cUkXfF0cuHnWnNd) to see how we achieve this. -- **A superior alternative to Apache Kafka**: 100% compatible with Apache Kafka and does not lose any key features, but cheaper and better. + - High throughput: Leverage pre-fetching, batch processing, and parallel technologies to maximize the capabilities of cloud object storage. Refer to the [AutoMQ Performance White Paper](https://www.automq.com/docs/automq/benchmarks/automq-vs-apache-kafka-benchmarks-and-cost?utm_source=github_automq) to see how we achieve this. + - Low Latency: AutoMQ defaults to running on S3 directly, resulting in hundreds of milliseconds of latency. The enterprise version offers single-digit millisecond latency. [Contact us](https://www.automq.com/contact?utm_source=github_automq) for more details. +- **Built-in Metrics Export**: Natively export Prometheus and OpenTelemetry metrics, supporting both push and pull. Ditch inefficient JMX and monitor your cluster with modern tools. Refer to [full metrics list](https://www.automq.com/docs/automq/observability/metrics?utm_source=github_automq) provided by AutoMQ. +- **100% Kafka Compatible**: Fully compatible with Apache Kafka, offering all features with greater cost-effectiveness and operational efficiency. ## ✨Architecture +AutoMQ is a fork of the open-source [Apache Kafka](https://github.com/apache/kafka). We've introduced a new storage engine based on object storage, transforming the classic shared-nothing architecture into a shared storage architecture. -![image](./docs/images/automq_vs_kafka.gif) - -AutoMQ adopts a Shared-Storage architecture, replacing the storage layer of Apache Kafka with a shared streaming storage library called [S3Stream](https://github.com/AutoMQ/automq/tree/main/s3stream) in a storage-compute separation manner, making the Broker completely stateless. +![image](./docs/images/automq_simple_arch.png) -Compared to the classic Kafka Shared-Nothing or Tiered-Storage architectures, AutoMQ's computing layer (Broker) is truly stateless, enabling features such as Auto-Scaling, Self-Balancing, and Partition Reassignment in Seconds that significantly reduce costs and improve efficiency. +Regarding the architecture of AutoMQ, it is fundamentally different from Kafka. The core difference lies in the storage layer of Apache Kafka and how we leverage object storage to achieve a stateless broker architecture. AutoMQ consists of below key components: +- S3 Storage Adapter: an adapter layer that reimplements the UnifiedLog, LocalLog, and LogSegment classes to create logs on S3 instead of a local disk. Traditional local disk storage is still supported if desired. +- S3Stream: a shared streaming storage library that encapsulates various storage modules, including WAL and object storage. WAL is a write-ahead log optimized for frequent writes and low IOPS to reduce S3 API costs. To boost read performance, we use LogCache and BlockCache for improved efficiency. +- Auto Balancer: a component that automatically balances traffic and partitions between brokers, eliminating the need for manual reassignment. Unlike Kafka, this built-in feature removes the need for cruise control. +- Rack-aware Router: Kafka has long faced cross-AZ traffic fees on AWS and GCP. Our shared storage architecture addresses this by using a rack-aware router to provide clients in different AZs with specific partition metadata, avoiding cross-AZ fees while exchanging data through object storage. -## ⛄ Get started with AutoMQ +For more on AutoMQ's architecture, visit [AutoMQ Architecture](https://www.automq.com/docs/automq/architecture/overview?utm_source=github_automq) or explore the source code directly. -### Deploy Locally on a Single Host -``` -curl https://download.automq.com/community_edition/standalone_deployment/install_run.sh | bash -``` - -The easiest way to run AutoMQ. You can experience features like **Partition Reassignment in Seconds** and **Continuous Self-Balancing** in your local machine. [Learn more](https://docs.automq.com/docs/automq-opensource/EsUBwQei4ilCDjkWb8WcbOZInwc) - -There are more deployment options available: -- [Deploy on Linux with 5 Nodes](https://docs.automq.com/docs/automq-opensource/IyXrw3lHriVPdQkQLDvcPGQdnNh) -- [Deploy on Kubernetes(Enterprise Edition Only)](https://docs.automq.com/docs/automq-opensource/KJtLwvdaPi7oznkX3lkcCR7fnte) -- [Runs on Ceph / MinIO / CubeFS / HDFS](https://docs.automq.com/docs/automq-opensource/RexrwfhKuiGChfk237QcEBIwnND) -- [Try AutoMQ on Alibaba Cloud Marketplace](https://market.aliyun.com/products/55530001/cmgj00065841.html) -- [Try AutoMQ on AWS Marketplace](https://docs.automq.com/automq-cloud/getting-started/install-byoc-environment/aws/install-env-from-marketplace) +## 🌟 Stay Ahead +Star AutoMQ on GitHub for instant updates on new releases. +![star-automq](https://github.com/user-attachments/assets/80a12561-2507-4283-8322-3512fec66f12) ## 💬 Community You can join the following groups or channels to discuss or ask questions about AutoMQ: - Ask questions or report a bug by [GitHub Issues](https://github.com/AutoMQ/automq/issues) -- Discuss about AutoMQ or Kafka by [Slack](https://join.slack.com/t/automq/shared_invite/zt-29h17vye9-thf31ebIVL9oXuRdACnOIA) or [Wechat Group](docs/images/automq-wechat.png) +- Discuss about AutoMQ or Kafka by [Slack](https://go.automq.com/slack) or [Wechat Group](docs/images/automq-wechat.png) ## 👥 How to contribute If you've found a problem with AutoMQ, please open a [GitHub Issues](https://github.com/AutoMQ/automq/issues). To contribute to AutoMQ please see [Code of Conduct](CODE_OF_CONDUCT.md) and [Contributing Guide](CONTRIBUTING_GUIDE.md). -We have a list of [good first issues](https://github.com/AutoMQ/automq/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) that help you to get started, gain experience, and get familiar with our contribution process. To claim one, simply reply with 'pick up' in the issue and the AutoMQ maintainers will assign the issue to you. If you have any questions about the 'good first issue' please feel free to ask. We will do our best to clarify any doubts you may have. - -## 👍 AutoMQ Business Edition -The business edition of AutoMQ provides a powerful and easy-to-use control plane to help you manage clusters effortlessly. Meanwhile, the control plane is more powerful in terms of availability and observability compared to the community edition. - -> You can check the difference between the community and business editions [here](https://www.automq.com/product). - - -Watch the following video and refer to our [docs](https://docs.automq.com/automq-cloud/getting-started/install-byoc-environment/aws/install-env-via-terraform-module) to see how to deploy AutoMQ Business Edition with 2 weeks free license for PoC. - - ⬇️ ⬇️ ⬇️ - -[![Deploy AutoMQ Business Edition with Terraform](https://img.youtube.com/vi/O40zp81x97w/0.jpg)](https://www.youtube.com/watch?v=O40zp81x97w) - - - -### Free trial of AutoMQ Business Edition -To allow users to experience the capabilities of the AutoMQ business edition without any barriers, click [here](https://www.automq.com/quick-start#Cloud?utm_source=github_automq_cloud) to apply for a no-obligation cluster trial, and note `AutoMQ Cloud Free Trial` in the message input box. We will immediately initialize an AutoMQ Cloud control panel for you soon in the cloud and give you the address of the control panel. Then, you can use the control panel to create a AutoMQ cluster or perform operations like scale in/out. - -No need to bind a credit card, no cost at all. We look forward to receiving valuable feedback from you to make our product better. If you want to proceed with a formal POC, you can also contact us through [Contact Us](https://automq66.feishu.cn/share/base/form/shrcnoqxslhYkujx6ULiMxOqkGh?utm_source=github_poc). We will further support your official POC. +We have a list of [good first issues](https://github.com/AutoMQ/automq/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) that help you to get started, gain experience, and get familiar with our contribution process. -## 👥 AutoMQ Use Cases -> Here are some of the users who have deployed AutoMQ in their production environments. +## 👍 AutoMQ Enterprise Edition +The enterprise edition of AutoMQ offers a robust, user-friendly control plane for seamless cluster management, with enhanced availability and observability over the open-source version. Additionally, we offer [Kafka Linking](https://www.automq.com/solutions/kafka-linking?utm_source=github_automq) for zero-downtime migration from any Kafka-compatible cluster to AutoMQ. -[![Click on the link to learn more about AutoMQ user cases.](./docs/images/customer.jpeg "AutoMQ Uses Case")](https://www.automq.com/customer) +[Contact us](https://www.automq.com/contact?utm_source=github_automq) for more information about the AutoMQ enterprise edition, and we'll gladly assist with your free trial. -## 🐱 The relationship with Apache Kafka +## 📜 License +AutoMQ is under the Apache 2.0 license. See the [LICENSE](https://github.com/AutoMQ/automq/blob/main/LICENSE) file for details. -AutoMQ is a fork of the open-source [Apache Kafka](https://github.com/apache/kafka). Based on the Apache Kafka codebase, we found an aspect at the LogSegment level, and replaced Kafka's storage layer with our self-developed cloud-first stream storage engine, [S3Stream](https://github.com/AutoMQ/automq/tree/main/s3stream). This engine can provide customers with high-performance, low-cost, and unlimited stream storage capabilities based on cloud storage like EBS WAL and S3. As such, AutoMQ completely retains the code of Kafka's computing layer and is 100% fully compatible with Apache Kafka. We appreciate the work done by the Apache Kafka community and will continue to embrace the Kafka community. +## 📝 Trademarks +Apache®, Apache Kafka®, Kafka®, Apache Iceberg®, Iceberg® and associated open source project names are trademarks of the Apache Software Foundation -## 🙋 Contact Us -Want to learn more, [Talk with our product experts](https://automq66.feishu.cn/share/base/form/shrcnoqxslhYkujx6ULiMxOqkGh). diff --git a/automq-log-uploader/README.md b/automq-log-uploader/README.md new file mode 100644 index 0000000000..8bc429ba98 --- /dev/null +++ b/automq-log-uploader/README.md @@ -0,0 +1,125 @@ +# AutoMQ Log Uploader Module + +This module provides asynchronous S3 log upload capability based on Log4j 1.x. Other submodules only need to depend on this module and configure it simply to synchronize logs to object storage. Core components: + +- `com.automq.log.S3RollingFileAppender`: Extends `RollingFileAppender`, pushes log events to the uploader while writing to local files. +- `com.automq.log.uploader.LogUploader`: Asynchronously buffers, compresses, and uploads logs; supports configuration switches and periodic cleanup. +- `com.automq.log.uploader.S3LogConfig`: Interface that abstracts the configuration required for uploading. Implementations must provide cluster ID, node ID, object storage instance, and leadership status. + +## Quick Integration + +1. Add dependency in your module's `build.gradle`: + ```groovy + implementation project(':automq-log-uploader') + ``` +2. Implement or provide an `S3LogConfig` instance and configure the appender: + + ```java + // Set up the S3LogConfig through your application + S3LogConfig config = // your S3LogConfig implementation + S3RollingFileAppender.setup(config); + ``` +3. Reference the Appender in `log4j.properties`: + + ```properties + log4j.appender.s3_uploader=com.automq.log.S3RollingFileAppender + log4j.appender.s3_uploader.File=logs/server.log + log4j.appender.s3_uploader.MaxFileSize=100MB + log4j.appender.s3_uploader.MaxBackupIndex=10 + log4j.appender.s3_uploader.layout=org.apache.log4j.PatternLayout + log4j.appender.s3_uploader.layout.ConversionPattern=[%d] %p %m (%c)%n + ``` + +## S3LogConfig Interface + +The `S3LogConfig` interface provides the configuration needed for log uploading: + +```java +public interface S3LogConfig { + boolean isEnabled(); // Whether S3 upload is enabled + String clusterId(); // Cluster identifier + int nodeId(); // Node identifier + ObjectStorage objectStorage(); // S3 object storage instance + boolean isLeader(); // Whether this node should upload logs +} +``` + + +The upload schedule can be overridden by environment variables: + +- `AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL`: Maximum upload interval (milliseconds). +- `AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL`: Retention period (milliseconds), old objects earlier than this time will be cleaned up. + +## Implementation Notes + +### Leader Selection + +The log uploader relies on the `S3LogConfig.isLeader()` method to determine whether the current node should upload logs and perform cleanup tasks. This avoids multiple nodes in a cluster simultaneously executing these operations. + +### Object Storage Path + +Logs are uploaded to object storage following this path pattern: +``` +automq/logs/{clusterId}/{nodeId}/{hour}/{uuid} +``` + +Where: +- `clusterId` and `nodeId` come from the S3LogConfig +- `hour` is the timestamp hour for log organization +- `uuid` is a unique identifier for each log batch + +## Usage Example + +Complete example of using the log uploader: + +```java +import com.automq.log.S3RollingFileAppender; +import com.automq.log.uploader.S3LogConfig; +import com.automq.stream.s3.operator.ObjectStorage; + +// Implement S3LogConfig +public class MyS3LogConfig implements S3LogConfig { + @Override + public boolean isEnabled() { + return true; // Enable S3 upload + } + + @Override + public String clusterId() { + return "my-cluster"; + } + + @Override + public int nodeId() { + return 1; + } + + @Override + public ObjectStorage objectStorage() { + // Return your ObjectStorage instance + return myObjectStorage; + } + + @Override + public boolean isLeader() { + // Return true if this node should upload logs + return isCurrentNodeLeader(); + } +} + +// Setup and use +S3LogConfig config = new MyS3LogConfig(); +S3RollingFileAppender.setup(config); + +// Configure Log4j to use the appender +// The appender will now automatically upload logs to S3 +``` + +## Lifecycle Management + +Remember to properly shutdown the log uploader when your application terminates: + +```java +// During application shutdown +S3RollingFileAppender.shutdown(); +``` diff --git a/automq-log-uploader/src/main/java/com/automq/log/S3RollingFileAppender.java b/automq-log-uploader/src/main/java/com/automq/log/S3RollingFileAppender.java new file mode 100644 index 0000000000..6e598df194 --- /dev/null +++ b/automq-log-uploader/src/main/java/com/automq/log/S3RollingFileAppender.java @@ -0,0 +1,105 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.log; + +import com.automq.log.uploader.LogRecorder; +import com.automq.log.uploader.LogUploader; +import com.automq.log.uploader.S3LogConfig; + +import org.apache.log4j.RollingFileAppender; +import org.apache.log4j.spi.LoggingEvent; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class S3RollingFileAppender extends RollingFileAppender { + + private static final Logger LOGGER = LoggerFactory.getLogger(S3RollingFileAppender.class); + private static final Object INIT_LOCK = new Object(); + + private static volatile LogUploader logUploaderInstance; + private static volatile S3LogConfig s3LogConfig; + + public S3RollingFileAppender() { + super(); + } + + public static void setup(S3LogConfig config) { + s3LogConfig = config; + synchronized (INIT_LOCK) { + if (logUploaderInstance != null) { + return; + } + try { + if (s3LogConfig == null) { + LOGGER.error("No s3LogConfig available; S3 log upload remains disabled."); + throw new RuntimeException("S3 log configuration is missing."); + } + if (!s3LogConfig.isEnabled() || s3LogConfig.objectStorage() == null) { + LOGGER.warn("S3 log upload is disabled by configuration."); + return; + } + + LogUploader uploader = new LogUploader(); + uploader.start(s3LogConfig); + logUploaderInstance = uploader; + LOGGER.info("S3RollingFileAppender initialized successfully using s3LogConfig {}.", s3LogConfig.getClass().getName()); + } catch (Exception e) { + LOGGER.error("Failed to initialize S3RollingFileAppender", e); + throw e; + } + } + } + + public static void shutdown() { + if (logUploaderInstance != null) { + synchronized (INIT_LOCK) { + if (logUploaderInstance != null) { + try { + logUploaderInstance.close(); + logUploaderInstance = null; + LOGGER.info("S3RollingFileAppender log uploader closed successfully."); + } catch (Exception e) { + LOGGER.error("Failed to close S3RollingFileAppender log uploader", e); + } + } + } + } + } + + @Override + protected void subAppend(LoggingEvent event) { + super.subAppend(event); + if (!closed && logUploaderInstance != null) { + LogRecorder.LogEvent logEvent = new LogRecorder.LogEvent( + event.getTimeStamp(), + event.getLevel().toString(), + event.getLoggerName(), + event.getRenderedMessage(), + event.getThrowableStrRep()); + + try { + logEvent.validate(); + logUploaderInstance.append(logEvent); + } catch (IllegalArgumentException e) { + errorHandler.error("Failed to validate and append log event", e, 0); + } + } + } +} diff --git a/automq-shell/src/main/java/com/automq/shell/log/LogRecorder.java b/automq-log-uploader/src/main/java/com/automq/log/uploader/LogRecorder.java similarity index 59% rename from automq-shell/src/main/java/com/automq/shell/log/LogRecorder.java rename to automq-log-uploader/src/main/java/com/automq/log/uploader/LogRecorder.java index e29c6e7274..04dc3e6914 100644 --- a/automq-shell/src/main/java/com/automq/shell/log/LogRecorder.java +++ b/automq-log-uploader/src/main/java/com/automq/log/uploader/LogRecorder.java @@ -1,15 +1,23 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package com.automq.shell.log; +package com.automq.log.uploader; import org.apache.commons.lang3.StringUtils; @@ -39,10 +47,10 @@ public void validate() { throw new IllegalArgumentException("Level cannot be blank"); } if (StringUtils.isBlank(logger)) { - throw new IllegalArgumentException("Level cannot be blank"); + throw new IllegalArgumentException("Logger cannot be blank"); } if (StringUtils.isBlank(message)) { - throw new IllegalArgumentException("Level cannot be blank"); + throw new IllegalArgumentException("Message cannot be blank"); } } diff --git a/automq-shell/src/main/java/com/automq/shell/log/LogUploader.java b/automq-log-uploader/src/main/java/com/automq/log/uploader/LogUploader.java similarity index 65% rename from automq-shell/src/main/java/com/automq/shell/log/LogUploader.java rename to automq-log-uploader/src/main/java/com/automq/log/uploader/LogUploader.java index 95aec2bd8e..01c6e4d9dd 100644 --- a/automq-shell/src/main/java/com/automq/shell/log/LogUploader.java +++ b/automq-log-uploader/src/main/java/com/automq/log/uploader/LogUploader.java @@ -1,17 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package com.automq.shell.log; +package com.automq.log.uploader; -import com.automq.shell.AutoMQApplication; +import com.automq.log.uploader.util.Utils; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorage.ObjectInfo; import com.automq.stream.s3.operator.ObjectStorage.ObjectPath; @@ -46,12 +54,14 @@ public class LogUploader implements LogRecorder { public static final int DEFAULT_MAX_QUEUE_SIZE = 64 * 1024; public static final int DEFAULT_BUFFER_SIZE = 16 * 1024 * 1024; - public static final int UPLOAD_INTERVAL = System.getenv("AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL") != null ? Integer.parseInt(System.getenv("AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL")) : 60 * 1000; - public static final int CLEANUP_INTERVAL = System.getenv("AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL") != null ? Integer.parseInt(System.getenv("AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL")) : 2 * 60 * 1000; + public static final int UPLOAD_INTERVAL = System.getenv("AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL") != null + ? Integer.parseInt(System.getenv("AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL")) + : 60 * 1000; + public static final int CLEANUP_INTERVAL = System.getenv("AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL") != null + ? Integer.parseInt(System.getenv("AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL")) + : 2 * 60 * 1000; public static final int MAX_JITTER_INTERVAL = 60 * 1000; - private static final LogUploader INSTANCE = new LogUploader(); - private final BlockingQueue queue = new LinkedBlockingQueue<>(DEFAULT_MAX_QUEUE_SIZE); private final ByteBuf uploadBuffer = Unpooled.directBuffer(DEFAULT_BUFFER_SIZE); private final Random random = new Random(); @@ -62,16 +72,42 @@ public class LogUploader implements LogRecorder { private volatile S3LogConfig config; - private volatile CompletableFuture startFuture; private ObjectStorage objectStorage; private Thread uploadThread; private Thread cleanupThread; - private LogUploader() { + public LogUploader() { } - public static LogUploader getInstance() { - return INSTANCE; + public synchronized void start(S3LogConfig config) { + if (this.config != null) { + LOGGER.warn("LogUploader is already started."); + return; + } + this.config = config; + if (!config.isEnabled() || config.objectStorage() == null) { + LOGGER.warn("LogUploader is disabled due to configuration."); + closed = true; + return; + } + + try { + this.objectStorage = config.objectStorage(); + this.uploadThread = new Thread(new UploadTask()); + this.uploadThread.setName("log-uploader-upload-thread"); + this.uploadThread.setDaemon(true); + this.uploadThread.start(); + + this.cleanupThread = new Thread(new CleanupTask()); + this.cleanupThread.setName("log-uploader-cleanup-thread"); + this.cleanupThread.setDaemon(true); + this.cleanupThread.start(); + + LOGGER.info("LogUploader started successfully."); + } catch (Exception e) { + LOGGER.error("Failed to start LogUploader", e); + closed = true; + } } public void close() throws InterruptedException { @@ -88,63 +124,15 @@ public void close() throws InterruptedException { @Override public boolean append(LogEvent event) { - if (!closed && couldUpload()) { + if (!closed) { return queue.offer(event); } return false; } - private boolean couldUpload() { - initConfiguration(); - boolean enabled = config != null && config.isEnabled() && config.objectStorage() != null; - - if (enabled) { - initUploadComponent(); - } - - return enabled && startFuture != null && startFuture.isDone(); - } - - private void initConfiguration() { - if (config == null) { - synchronized (this) { - if (config == null) { - config = AutoMQApplication.getBean(S3LogConfig.class); - } - } - } - } - - private void initUploadComponent() { - if (startFuture == null) { - synchronized (this) { - if (startFuture == null) { - startFuture = CompletableFuture.runAsync(() -> { - try { - objectStorage = config.objectStorage(); - uploadThread = new Thread(new UploadTask()); - uploadThread.setName("log-uploader-upload-thread"); - uploadThread.setDaemon(true); - uploadThread.start(); - - cleanupThread = new Thread(new CleanupTask()); - cleanupThread.setName("log-uploader-cleanup-thread"); - cleanupThread.setDaemon(true); - cleanupThread.start(); - - startFuture.complete(null); - } catch (Exception e) { - LOGGER.error("Initialize log uploader failed", e); - } - }, command -> new Thread(command).start()); - } - } - } - } - private class UploadTask implements Runnable { - public String formatTimestampInMillis(long timestamp) { + private String formatTimestampInMillis(long timestamp) { return ZonedDateTime.ofInstant(Instant.ofEpochMilli(timestamp), ZoneId.systemDefault()) .format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS Z")); } @@ -156,7 +144,6 @@ public void run() { long now = System.currentTimeMillis(); LogEvent event = queue.poll(1, TimeUnit.SECONDS); if (event != null) { - // DateTime Level [Logger] Message \n stackTrace StringBuilder logLine = new StringBuilder() .append(formatTimestampInMillis(event.timestampMillis())) .append(" ") @@ -195,25 +182,22 @@ public void run() { private void upload(long now) { if (uploadBuffer.readableBytes() > 0) { - if (couldUpload()) { - try { - while (!Thread.currentThread().isInterrupted()) { - if (objectStorage == null) { - break; - } - - try { - String objectKey = getObjectKey(); - objectStorage.write(WriteOptions.DEFAULT, objectKey, uploadBuffer.retainedSlice().asReadOnly()).get(); - break; - } catch (Exception e) { - e.printStackTrace(System.err); - Thread.sleep(1000); - } + try { + while (!Thread.currentThread().isInterrupted()) { + if (objectStorage == null) { + break; + } + try { + String objectKey = getObjectKey(); + objectStorage.write(WriteOptions.DEFAULT, objectKey, Utils.compress(uploadBuffer.slice().asReadOnly())).get(); + break; + } catch (Exception e) { + LOGGER.warn("Failed to upload logs, will retry", e); + Thread.sleep(1000); } - } catch (InterruptedException e) { - //ignore } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); } uploadBuffer.clear(); lastUploadTimestamp = now; @@ -228,12 +212,11 @@ private class CleanupTask implements Runnable { public void run() { while (!Thread.currentThread().isInterrupted()) { try { - if (closed || !config.isActiveController()) { + if (closed || !config.isLeader()) { Thread.sleep(Duration.ofMinutes(1).toMillis()); continue; } long expiredTime = System.currentTimeMillis() - CLEANUP_INTERVAL; - List objects = objectStorage.list(String.format("automq/logs/%s", config.clusterId())).join(); if (!objects.isEmpty()) { @@ -243,7 +226,6 @@ public void run() { .collect(Collectors.toList()); if (!keyList.isEmpty()) { - // Some of s3 implements allow only 1000 keys per request. CompletableFuture[] deleteFutures = Lists.partition(keyList, 1000) .stream() .map(objectStorage::delete) @@ -251,7 +233,6 @@ public void run() { CompletableFuture.allOf(deleteFutures).join(); } } - Thread.sleep(Duration.ofMinutes(1).toMillis()); } catch (InterruptedException e) { break; @@ -266,5 +247,4 @@ private String getObjectKey() { String hour = LocalDateTime.now(ZoneOffset.UTC).format(DateTimeFormatter.ofPattern("yyyyMMddHH")); return String.format("automq/logs/%s/%s/%s/%s", config.clusterId(), config.nodeId(), hour, UUID.randomUUID()); } - } diff --git a/automq-log-uploader/src/main/java/com/automq/log/uploader/S3LogConfig.java b/automq-log-uploader/src/main/java/com/automq/log/uploader/S3LogConfig.java new file mode 100644 index 0000000000..f239d86ac8 --- /dev/null +++ b/automq-log-uploader/src/main/java/com/automq/log/uploader/S3LogConfig.java @@ -0,0 +1,34 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.log.uploader; + +import com.automq.stream.s3.operator.ObjectStorage; + +public interface S3LogConfig { + boolean isEnabled(); + + String clusterId(); + + int nodeId(); + + ObjectStorage objectStorage(); + + boolean isLeader(); +} diff --git a/automq-log-uploader/src/main/java/com/automq/log/uploader/util/Utils.java b/automq-log-uploader/src/main/java/com/automq/log/uploader/util/Utils.java new file mode 100644 index 0000000000..442d6aac84 --- /dev/null +++ b/automq-log-uploader/src/main/java/com/automq/log/uploader/util/Utils.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.log.uploader.util; + +import com.automq.stream.s3.ByteBufAlloc; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +import io.netty.buffer.ByteBuf; + +public class Utils { + + private Utils() { + } + + public static ByteBuf compress(ByteBuf input) throws IOException { + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream)) { + byte[] buffer = new byte[input.readableBytes()]; + input.readBytes(buffer); + gzipOutputStream.write(buffer); + } + + ByteBuf compressed = ByteBufAlloc.byteBuffer(byteArrayOutputStream.size()); + compressed.writeBytes(byteArrayOutputStream.toByteArray()); + return compressed; + } + + public static ByteBuf decompress(ByteBuf input) throws IOException { + byte[] compressedData = new byte[input.readableBytes()]; + input.readBytes(compressedData); + ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(compressedData); + + try (GZIPInputStream gzipInputStream = new GZIPInputStream(byteArrayInputStream); + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream()) { + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = gzipInputStream.read(buffer)) != -1) { + byteArrayOutputStream.write(buffer, 0, bytesRead); + } + + byte[] uncompressedData = byteArrayOutputStream.toByteArray(); + ByteBuf output = ByteBufAlloc.byteBuffer(uncompressedData.length); + output.writeBytes(uncompressedData); + return output; + } + } +} diff --git a/automq-metrics/README.md b/automq-metrics/README.md new file mode 100644 index 0000000000..0ef70680ae --- /dev/null +++ b/automq-metrics/README.md @@ -0,0 +1,459 @@ +# AutoMQ automq-metrics Module + +## Module Structure + +``` +com.automq.opentelemetry/ +├── AutoMQTelemetryManager.java # Main management class for initialization and lifecycle +├── TelemetryConstants.java # Constants definition +├── common/ +│ ├── OTLPCompressionType.java # OTLP compression types +│ └── OTLPProtocol.java # OTLP protocol types +├── exporter/ +│ ├── MetricsExporter.java # Exporter interface +│ ├── MetricsExportConfig.java # Export configuration +│ ├── MetricsExporterProvider.java # Exporter factory provider +│ ├── MetricsExporterType.java # Exporter type enumeration +│ ├── MetricsExporterURI.java # URI parser for exporters +│ ├── OTLPMetricsExporter.java # OTLP exporter implementation +│ ├── PrometheusMetricsExporter.java # Prometheus exporter implementation +│ └── s3/ # S3 metrics exporter implementation +│ ├── CompressionUtils.java # Utility for data compression +│ ├── PrometheusUtils.java # Utilities for Prometheus format +│ ├── S3MetricsExporter.java # S3 metrics exporter implementation +│ └── S3MetricsExporterAdapter.java # Adapter to handle S3 metrics export +└── yammer/ + ├── DeltaHistogram.java # Delta histogram implementation + ├── OTelMetricUtils.java # OpenTelemetry metrics utilities + ├── YammerMetricsProcessor.java # Yammer metrics processor + └── YammerMetricsReporter.java # Yammer metrics reporter +``` + +The AutoMQ OpenTelemetry module is a telemetry data collection and export component based on OpenTelemetry SDK, specifically designed for AutoMQ Kafka. This module provides unified telemetry data management capabilities, supporting the collection of JVM metrics, JMX metrics, and Yammer metrics, and can export data to Prometheus, OTLP-compatible backend systems, or S3-compatible storage. + +## Core Features + +### 1. Metrics Collection +- **JVM Metrics**: Automatically collect JVM runtime metrics including CPU, memory pools, garbage collection, threads, etc. +- **JMX Metrics**: Define and collect JMX Bean metrics through configuration files +- **Yammer Metrics**: Bridge existing Kafka Yammer metrics system to OpenTelemetry + +### 2. Multiple Exporter Support +- **Prometheus**: Expose metrics in Prometheus format through HTTP server +- **OTLP**: Support both gRPC and HTTP/Protobuf protocols for exporting to OTLP backends +- **S3**: Export metrics to S3-compatible object storage systems + +### 3. Flexible Configuration +- Support parameter settings through Properties configuration files +- Configurable export intervals, compression methods, timeout values, etc. +- Support metric cardinality limits to control memory usage + +## Module Structure + +``` +com.automq.opentelemetry/ +├── AutoMQTelemetryManager.java # Main management class for initialization and lifecycle +├── TelemetryConfig.java # Configuration management class +├── TelemetryConstants.java # Constants definition +├── common/ +│ └── MetricsUtils.java # Metrics utility class +├── exporter/ +│ ├── MetricsExporter.java # Exporter interface +│ ├── MetricsExporterURI.java # URI parser +│ ���── OTLPMetricsExporter.java # OTLP exporter implementation +│ ├── PrometheusMetricsExporter.java # Prometheus exporter implementation +│ └── s3/ # S3 metrics exporter implementation +│ ├── CompressionUtils.java # Utility for data compression +│ ├── PrometheusUtils.java # Utilities for Prometheus format +│ ├── S3MetricsConfig.java # Configuration interface +│ ├── S3MetricsExporter.java # S3 metrics exporter implementation +│ ├── S3MetricsExporterAdapter.java # Adapter to handle S3 metrics export +│ ├── LeaderNodeSelector.java # Interface for node selection logic +│ └── LeaderNodeSelectors.java # Factory for node selector implementations +└── yammer/ + ├── DeltaHistogram.java # Delta histogram implementation + ├── OTelMetricUtils.java # OpenTelemetry metrics utilities + ├── YammerMetricsProcessor.java # Yammer metrics processor + └── YammerMetricsReporter.java # Yammer metrics reporter +``` + +## Quick Start + +### 1. Basic Usage + +```java +import com.automq.opentelemetry.AutoMQTelemetryManager; +import com.automq.opentelemetry.exporter.MetricsExportConfig; + +// Implement MetricsExportConfig +public class MyMetricsExportConfig implements MetricsExportConfig { + @Override + public String clusterId() { return "my-cluster"; } + + @Override + public boolean isLeader() { return true; } + + @Override + public int nodeId() { return 1; } + + @Override + public ObjectStorage objectStorage() { + // Return your object storage instance for S3 exports + return myObjectStorage; + } + + @Override + public List> baseLabels() { + return Arrays.asList( + Pair.of("environment", "production"), + Pair.of("region", "us-east-1") + ); + } + + @Override + public int intervalMs() { return 60000; } // 60 seconds +} + +// Create export configuration +MetricsExportConfig config = new MyMetricsExportConfig(); + +// Initialize telemetry manager singleton +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "prometheus://localhost:9090", // exporter URI + "automq-kafka", // service name + "broker-1", // instance ID + config // export config +); + +// Start Yammer metrics reporting (optional) +MetricsRegistry yammerRegistry = // Get Kafka's Yammer registry +manager.startYammerMetricsReporter(yammerRegistry); + +// Application running... + +// Shutdown telemetry system +AutoMQTelemetryManager.shutdownInstance(); +``` + +### 2. Get Meter Instance + +```java +// Get the singleton instance +AutoMQTelemetryManager manager = AutoMQTelemetryManager.getInstance(); + +// Get Meter for custom metrics +Meter meter = manager.getMeter(); + +// Create custom metrics +LongCounter requestCounter = meter + .counterBuilder("http_requests_total") + .setDescription("Total number of HTTP requests") + .build(); + +requestCounter.add(1, Attributes.of(AttributeKey.stringKey("method"), "GET")); +``` + +## Configuration + +### Basic Configuration + +Configuration is provided through the `MetricsExportConfig` interface and constructor parameters: + +| Parameter | Description | Example | +|-----------|-------------|---------| +| `exporterUri` | Metrics exporter URI | `prometheus://localhost:9090` | +| `serviceName` | Service name for telemetry | `automq-kafka` | +| `instanceId` | Unique service instance ID | `broker-1` | +| `config` | MetricsExportConfig implementation | See example above | + +### Exporter Configuration + +All configuration is done through the `MetricsExportConfig` interface and constructor parameters. Export intervals, compression settings, and other options are controlled through: + +1. **Exporter URI**: Determines the export destination and protocol +2. **MetricsExportConfig**: Provides cluster information, intervals, and base labels +3. **Constructor parameters**: Service name and instance ID + +#### Prometheus Exporter +```java +// Use prometheus:// URI scheme +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "prometheus://localhost:9090", + "automq-kafka", + "broker-1", + config +); +``` + +#### OTLP Exporter +```java +// Use otlp:// URI scheme with optional query parameters +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "otlp://localhost:4317?protocol=grpc&compression=gzip&timeout=30000", + "automq-kafka", + "broker-1", + config +); +``` + +#### S3 Metrics Exporter +```java +// Use s3:// URI scheme +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "s3://access-key:secret-key@my-bucket.s3.amazonaws.com", + "automq-kafka", + "broker-1", + config // config.clusterId(), nodeId(), isLeader() used for S3 export +); +``` + +Example usage with S3 exporter: + +```java +// Implementation for S3 export configuration +public class S3MetricsExportConfig implements MetricsExportConfig { + private final ObjectStorage objectStorage; + + public S3MetricsExportConfig(ObjectStorage objectStorage) { + this.objectStorage = objectStorage; + } + + @Override + public String clusterId() { return "my-kafka-cluster"; } + + @Override + public boolean isLeader() { + // Only one node in the cluster should return true + return isCurrentNodeLeader(); + } + + @Override + public int nodeId() { return 1; } + + @Override + public ObjectStorage objectStorage() { return objectStorage; } + + @Override + public List> baseLabels() { + return Arrays.asList(Pair.of("environment", "production")); + } + + @Override + public int intervalMs() { return 60000; } +} + +// Initialize telemetry manager with S3 export +ObjectStorage objectStorage = // Create your object storage instance +MetricsExportConfig config = new S3MetricsExportConfig(objectStorage); + +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "s3://access-key:secret-key@my-bucket.s3.amazonaws.com", + "automq-kafka", + "broker-1", + config +); + +// Application running... + +// Shutdown telemetry system +AutoMQTelemetryManager.shutdownInstance(); +``` + +### JMX Metrics Configuration + +Define JMX metrics collection rules through YAML configuration files: + +```java +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + exporterUri, serviceName, instanceId, config +); + +// Set JMX config paths after initialization +manager.setJmxConfigPaths("/jmx-config.yaml,/kafka-jmx.yaml"); +``` + +#### Configuration File Requirements + +1. **Directory Requirements**: + - Configuration files must be placed in the project's classpath (e.g., `src/main/resources` directory) + - Support subdirectory structure, e.g., `/config/jmx-metrics.yaml` + +2. **Path Format**: + - Paths must start with `/` to indicate starting from classpath root + - Multiple configuration files separated by commas + +3. **File Format**: + - Use YAML format (`.yaml` or `.yml` extension) + - Filenames can be customized, meaningful names are recommended + +#### Recommended Directory Structure + +``` +src/main/resources/ +├── jmx-kafka-broker.yaml # Kafka Broker metrics configuration +├── jmx-kafka-consumer.yaml # Kafka Consumer metrics configuration +├── jmx-kafka-producer.yaml # Kafka Producer metrics configuration +└── config/ + ├── custom-jmx.yaml # Custom JMX metrics configuration + └── third-party-jmx.yaml # Third-party component JMX configuration +``` + +JMX configuration file example (`jmx-config.yaml`): +```yaml +rules: + - bean: kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec + metricAttribute: + name: kafka_server_broker_topic_messages_in_per_sec + description: Messages in per second + unit: "1/s" + attributes: + - name: topic + value: topic +``` + +## Supported Metric Types + +### 1. JVM Metrics +- Memory usage (heap memory, non-heap memory, memory pools) +- CPU usage +- Garbage collection statistics +- Thread states + +### 2. Kafka Metrics +Through Yammer metrics bridging, supports the following types of Kafka metrics: +- `BytesInPerSec` - Bytes input per second +- `BytesOutPerSec` - Bytes output per second +- `Size` - Log size (for identifying idle partitions) + +### 3. Custom Metrics +Support creating custom metrics through OpenTelemetry API: +- Counter +- Gauge +- Histogram +- UpDownCounter + +## Best Practices + +### 1. Production Environment Configuration + +```java +public class ProductionMetricsConfig implements MetricsExportConfig { + @Override + public String clusterId() { return "production-cluster"; } + + @Override + public boolean isLeader() { + // Implement your leader election logic + return isCurrentNodeController(); + } + + @Override + public int nodeId() { return getCurrentNodeId(); } + + @Override + public ObjectStorage objectStorage() { + return productionObjectStorage; + } + + @Override + public List> baseLabels() { + return Arrays.asList( + Pair.of("environment", "production"), + Pair.of("region", System.getenv("AWS_REGION")), + Pair.of("version", getApplicationVersion()) + ); + } + + @Override + public int intervalMs() { return 60000; } // 1 minute +} + +// Initialize for production +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "prometheus://0.0.0.0:9090", // Or S3 URI for object storage export + "automq-kafka", + System.getenv("HOSTNAME"), + new ProductionMetricsConfig() +); +``` + +### 2. Development Environment Configuration + +```java +public class DevelopmentMetricsConfig implements MetricsExportConfig { + @Override + public String clusterId() { return "dev-cluster"; } + + @Override + public boolean isLeader() { return true; } // Single node in dev + + @Override + public int nodeId() { return 1; } + + @Override + public ObjectStorage objectStorage() { return null; } // Not needed for OTLP + + @Override + public List> baseLabels() { + return Arrays.asList(Pair.of("environment", "development")); + } + + @Override + public int intervalMs() { return 10000; } // 10 seconds for faster feedback +} + +// Initialize for development +AutoMQTelemetryManager manager = AutoMQTelemetryManager.initializeInstance( + "otlp://localhost:4317", + "automq-kafka-dev", + "local-dev", + new DevelopmentMetricsConfig() +); +``` + +### 3. Resource Management +- Set appropriate metric cardinality limits to avoid memory leaks +- Call `shutdown()` method when application closes to release resources +- Monitor exporter health status + +## Troubleshooting + +### Common Issues + +1. **Metrics not exported** + - Check if exporter URI passed to `initializeInstance()` is correct + - Verify target endpoint is reachable + - Check error messages in logs + - Ensure `MetricsExportConfig.intervalMs()` returns reasonable value + +2. **JMX metrics missing** + - Confirm JMX configuration file path set via `setJmxConfigPaths()` is correct + - Check YAML configuration file format + - Verify JMX Bean exists + - Ensure files are in classpath + +3. **High memory usage** + - Implement cardinality limits in your `MetricsExportConfig` + - Check for high cardinality labels in `baseLabels()` + - Consider increasing export interval via `intervalMs()` + +### Logging Configuration + +Enable debug logging for more information using your logging framework configuration (e.g., logback.xml, log4j2.xml): + +```xml + + + +``` + +## Dependencies + +- Java 8+ +- OpenTelemetry SDK 1.30+ +- Apache Commons Lang3 +- SLF4J logging framework + +## License + +This module is open source under the Apache License 2.0. diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/AutoMQTelemetryManager.java b/automq-metrics/src/main/java/com/automq/opentelemetry/AutoMQTelemetryManager.java new file mode 100644 index 0000000000..38beb126e3 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/AutoMQTelemetryManager.java @@ -0,0 +1,330 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry; + +import com.automq.opentelemetry.exporter.MetricsExportConfig; +import com.automq.opentelemetry.exporter.MetricsExporter; +import com.automq.opentelemetry.exporter.MetricsExporterURI; +import com.automq.opentelemetry.yammer.YammerMetricsReporter; +import com.yammer.metrics.core.MetricsRegistry; + +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.slf4j.bridge.SLF4JBridgeHandler; + +import java.io.IOException; +import java.io.InputStream; +import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.common.AttributesBuilder; +import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; +import io.opentelemetry.context.propagation.ContextPropagators; +import io.opentelemetry.context.propagation.TextMapPropagator; +import io.opentelemetry.instrumentation.jmx.engine.JmxMetricInsight; +import io.opentelemetry.instrumentation.jmx.engine.MetricConfiguration; +import io.opentelemetry.instrumentation.jmx.yaml.RuleParser; +import io.opentelemetry.instrumentation.runtimemetrics.java8.Cpu; +import io.opentelemetry.instrumentation.runtimemetrics.java8.GarbageCollector; +import io.opentelemetry.instrumentation.runtimemetrics.java8.MemoryPools; +import io.opentelemetry.instrumentation.runtimemetrics.java8.Threads; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.metrics.SdkMeterProvider; +import io.opentelemetry.sdk.metrics.SdkMeterProviderBuilder; +import io.opentelemetry.sdk.metrics.export.MetricReader; +import io.opentelemetry.sdk.metrics.internal.SdkMeterProviderUtil; +import io.opentelemetry.sdk.resources.Resource; + +/** + * The main manager for AutoMQ telemetry. + * This class is responsible for initializing, configuring, and managing the lifecycle of all + * telemetry components, including the OpenTelemetry SDK, metric exporters, and various metric sources. + */ +public class AutoMQTelemetryManager { + private static final Logger LOGGER = LoggerFactory.getLogger(AutoMQTelemetryManager.class); + + // Singleton instance support + private static volatile AutoMQTelemetryManager instance; + private static final Object LOCK = new Object(); + + private final String exporterUri; + private final String serviceName; + private final String instanceId; + private final MetricsExportConfig metricsExportConfig; + private final List metricReaders = new ArrayList<>(); + private final List autoCloseableList; + private OpenTelemetrySdk openTelemetrySdk; + private YammerMetricsReporter yammerReporter; + + private int metricCardinalityLimit = TelemetryConstants.DEFAULT_METRIC_CARDINALITY_LIMIT; + private String jmxConfigPath; + + /** + * Constructs a new Telemetry Manager with the given configuration. + * + * @param exporterUri The metrics exporter URI. + * @param serviceName The service name to be used in telemetry data. + * @param instanceId The unique instance ID for this service instance. + * @param metricsExportConfig The metrics configuration. + */ + public AutoMQTelemetryManager(String exporterUri, String serviceName, String instanceId, MetricsExportConfig metricsExportConfig) { + this.exporterUri = exporterUri; + this.serviceName = serviceName; + this.instanceId = instanceId; + this.metricsExportConfig = metricsExportConfig; + this.autoCloseableList = new ArrayList<>(); + // Redirect JUL from OpenTelemetry SDK to SLF4J for unified logging + SLF4JBridgeHandler.removeHandlersForRootLogger(); + SLF4JBridgeHandler.install(); + } + + /** + * Gets the singleton instance of AutoMQTelemetryManager. + * Returns null if no instance has been initialized. + * + * @return the singleton instance, or null if not initialized + */ + public static AutoMQTelemetryManager getInstance() { + return instance; + } + + /** + * Initializes the singleton instance with the given configuration. + * This method should be called before any other components try to access the instance. + * + * @param exporterUri The metrics exporter URI. + * @param serviceName The service name to be used in telemetry data. + * @param instanceId The unique instance ID for this service instance. + * @param metricsExportConfig The metrics configuration. + * @return the initialized singleton instance + */ + public static AutoMQTelemetryManager initializeInstance(String exporterUri, String serviceName, String instanceId, MetricsExportConfig metricsExportConfig) { + if (instance == null) { + synchronized (LOCK) { + if (instance == null) { + AutoMQTelemetryManager newInstance = new AutoMQTelemetryManager(exporterUri, serviceName, instanceId, metricsExportConfig); + newInstance.init(); + instance = newInstance; + LOGGER.info("AutoMQTelemetryManager singleton instance initialized"); + } + } + } + return instance; + } + + /** + * Shuts down the singleton instance and releases all resources. + */ + public static void shutdownInstance() { + if (instance != null) { + synchronized (LOCK) { + if (instance != null) { + instance.shutdown(); + instance = null; + LOGGER.info("AutoMQTelemetryManager singleton instance shutdown"); + } + } + } + } + + /** + * Initializes the telemetry system. This method sets up the OpenTelemetry SDK, + * configures exporters, and registers JVM and JMX metrics. + */ + public void init() { + SdkMeterProvider meterProvider = buildMeterProvider(); + + this.openTelemetrySdk = OpenTelemetrySdk.builder() + .setMeterProvider(meterProvider) + .setPropagators(ContextPropagators.create(TextMapPropagator.composite( + W3CTraceContextPropagator.getInstance(), W3CBaggagePropagator.getInstance()))) + .buildAndRegisterGlobal(); + + // Register JVM and JMX metrics + registerJvmMetrics(openTelemetrySdk); + registerJmxMetrics(openTelemetrySdk); + + LOGGER.info("AutoMQ Telemetry Manager initialized successfully."); + } + + private SdkMeterProvider buildMeterProvider() { + String hostName; + try { + hostName = InetAddress.getLocalHost().getHostName(); + } catch (UnknownHostException e) { + hostName = "unknown-host"; + } + AttributesBuilder attrsBuilder = Attributes.builder() + .put(TelemetryConstants.SERVICE_NAME_KEY, serviceName) + .put(TelemetryConstants.SERVICE_INSTANCE_ID_KEY, instanceId) + .put(TelemetryConstants.HOST_NAME_KEY, hostName) + // Add attributes for Prometheus compatibility + .put(TelemetryConstants.PROMETHEUS_JOB_KEY, serviceName) + .put(TelemetryConstants.PROMETHEUS_INSTANCE_KEY, instanceId); + + for (Pair label : metricsExportConfig.baseLabels()) { + attrsBuilder.put(label.getKey(), label.getValue()); + } + + Resource resource = Resource.getDefault().merge(Resource.create(attrsBuilder.build())); + SdkMeterProviderBuilder meterProviderBuilder = SdkMeterProvider.builder().setResource(resource); + + // Configure exporters from URI + MetricsExporterURI exporterURI = buildMetricsExporterURI(exporterUri, metricsExportConfig); + for (MetricsExporter exporter : exporterURI.getMetricsExporters()) { + MetricReader reader = exporter.asMetricReader(); + metricReaders.add(reader); + SdkMeterProviderUtil.registerMetricReaderWithCardinalitySelector(meterProviderBuilder, reader, + instrumentType -> metricCardinalityLimit); + } + + return meterProviderBuilder.build(); + } + + protected MetricsExporterURI buildMetricsExporterURI(String exporterUri, MetricsExportConfig metricsExportConfig) { + return MetricsExporterURI.parse(exporterUri, metricsExportConfig); + } + + private void registerJvmMetrics(OpenTelemetry openTelemetry) { + autoCloseableList.addAll(MemoryPools.registerObservers(openTelemetry)); + autoCloseableList.addAll(Cpu.registerObservers(openTelemetry)); + autoCloseableList.addAll(GarbageCollector.registerObservers(openTelemetry)); + autoCloseableList.addAll(Threads.registerObservers(openTelemetry)); + LOGGER.info("JVM metrics registered."); + } + + @SuppressWarnings({"NP_LOAD_OF_KNOWN_NULL_VALUE", "RCN_REDUNDANT_NULLCHECK_OF_NULL_VALUE"}) + private void registerJmxMetrics(OpenTelemetry openTelemetry) { + List jmxConfigPaths = getJmxConfigPaths(); + if (jmxConfigPaths.isEmpty()) { + LOGGER.info("No JMX metric config paths provided, skipping JMX metrics registration."); + return; + } + + JmxMetricInsight jmxMetricInsight = JmxMetricInsight.createService(openTelemetry, metricsExportConfig.intervalMs()); + MetricConfiguration metricConfig = new MetricConfiguration(); + + for (String path : jmxConfigPaths) { + try (InputStream ins = this.getClass().getResourceAsStream(path)) { + if (ins == null) { + LOGGER.error("JMX config file not found in classpath: {}", path); + continue; + } + RuleParser parser = RuleParser.get(); + parser.addMetricDefsTo(metricConfig, ins, path); + } catch (Exception e) { + LOGGER.error("Failed to parse JMX config file: {}", path, e); + } + } + + jmxMetricInsight.start(metricConfig); + // JmxMetricInsight doesn't implement Closeable, but we can create a wrapper + + LOGGER.info("JMX metrics registered with config paths: {}", jmxConfigPaths); + } + + public List getJmxConfigPaths() { + if (StringUtils.isEmpty(jmxConfigPath)) { + return Collections.emptyList(); + } + return Stream.of(jmxConfigPath.split(",")) + .map(String::trim) + .filter(s -> !s.isEmpty()) + .collect(Collectors.toList()); + } + + /** + * Starts reporting metrics from a given Yammer MetricsRegistry. + * + * @param registry The Yammer registry to bridge metrics from. + */ + public void startYammerMetricsReporter(MetricsRegistry registry) { + if (this.openTelemetrySdk == null) { + throw new IllegalStateException("TelemetryManager is not initialized. Call init() first."); + } + if (registry == null) { + LOGGER.warn("Yammer MetricsRegistry is null, skipping reporter start."); + return; + } + this.yammerReporter = new YammerMetricsReporter(registry); + this.yammerReporter.start(getMeter()); + } + + public void shutdown() { + autoCloseableList.forEach(autoCloseable -> { + try { + autoCloseable.close(); + } catch (Exception e) { + LOGGER.error("Failed to close auto closeable", e); + } + }); + metricReaders.forEach(metricReader -> { + metricReader.forceFlush(); + try { + metricReader.close(); + } catch (IOException e) { + LOGGER.error("Failed to close metric reader", e); + } + }); + if (openTelemetrySdk != null) { + openTelemetrySdk.close(); + } + } + + /** + * get YammerMetricsReporter instance. + * + * @return The YammerMetricsReporter instance. + */ + public YammerMetricsReporter getYammerReporter() { + return this.yammerReporter; + } + + public void setMetricCardinalityLimit(int limit) { + this.metricCardinalityLimit = limit; + } + + public void setJmxConfigPaths(String jmxConfigPaths) { + this.jmxConfigPath = jmxConfigPaths; + } + + /** + * Gets the default meter from the initialized OpenTelemetry SDK. + * + * @return The meter instance. + */ + public Meter getMeter() { + if (this.openTelemetrySdk == null) { + throw new IllegalStateException("TelemetryManager is not initialized. Call init() first."); + } + return this.openTelemetrySdk.getMeter(TelemetryConstants.TELEMETRY_SCOPE_NAME); + } +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/TelemetryConstants.java b/automq-metrics/src/main/java/com/automq/opentelemetry/TelemetryConstants.java new file mode 100644 index 0000000000..54392ba5db --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/TelemetryConstants.java @@ -0,0 +1,54 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry; + +import io.opentelemetry.api.common.AttributeKey; + +/** + * Constants for telemetry, including configuration keys, attribute keys, and default values. + */ +public class TelemetryConstants { + + //################################################################ + // Service and Resource Attributes + //################################################################ + public static final String SERVICE_NAME_KEY = "service.name"; + public static final String SERVICE_INSTANCE_ID_KEY = "service.instance.id"; + public static final String HOST_NAME_KEY = "host.name"; + public static final String TELEMETRY_SCOPE_NAME = "automq_for_kafka"; + + /** + * The cardinality limit for any single metric. + */ + public static final String METRIC_CARDINALITY_LIMIT_KEY = "automq.telemetry.metric.cardinality.limit"; + public static final int DEFAULT_METRIC_CARDINALITY_LIMIT = 20000; + + //################################################################ + // Prometheus specific Attributes, for compatibility + //################################################################ + public static final String PROMETHEUS_JOB_KEY = "job"; + public static final String PROMETHEUS_INSTANCE_KEY = "instance"; + + //################################################################ + // Custom Kafka-related Attribute Keys + //################################################################ + public static final AttributeKey START_OFFSET_KEY = AttributeKey.longKey("startOffset"); + public static final AttributeKey END_OFFSET_KEY = AttributeKey.longKey("endOffset"); +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/common/OTLPCompressionType.java b/automq-metrics/src/main/java/com/automq/opentelemetry/common/OTLPCompressionType.java new file mode 100644 index 0000000000..4833159149 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/common/OTLPCompressionType.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.common; + +public enum OTLPCompressionType { + GZIP("gzip"), + NONE("none"); + + private final String type; + + OTLPCompressionType(String type) { + this.type = type; + } + + public String getType() { + return type; + } + + public static OTLPCompressionType fromString(String type) { + for (OTLPCompressionType compressionType : OTLPCompressionType.values()) { + if (compressionType.getType().equalsIgnoreCase(type)) { + return compressionType; + } + } + throw new IllegalArgumentException("Invalid OTLP compression type: " + type); + } +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/common/OTLPProtocol.java b/automq-metrics/src/main/java/com/automq/opentelemetry/common/OTLPProtocol.java new file mode 100644 index 0000000000..69f3cd1918 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/common/OTLPProtocol.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.common; + +public enum OTLPProtocol { + GRPC("grpc"), + HTTP("http"); + + private final String protocol; + + OTLPProtocol(String protocol) { + this.protocol = protocol; + } + + public String getProtocol() { + return protocol; + } + + public static OTLPProtocol fromString(String protocol) { + for (OTLPProtocol otlpProtocol : OTLPProtocol.values()) { + if (otlpProtocol.getProtocol().equalsIgnoreCase(protocol)) { + return otlpProtocol; + } + } + throw new IllegalArgumentException("Invalid OTLP protocol: " + protocol); + } +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExportConfig.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExportConfig.java new file mode 100644 index 0000000000..3e54b2b172 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExportConfig.java @@ -0,0 +1,68 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter; + +import com.automq.stream.s3.operator.ObjectStorage; + +import org.apache.commons.lang3.tuple.Pair; + +import java.util.List; + +/** + * Configuration interface for metrics exporter. + */ +public interface MetricsExportConfig { + + /** + * Get the cluster ID. + * @return The cluster ID. + */ + String clusterId(); + + /** + * Check if the current node is a primary node for metrics upload. + * @return True if the current node should upload metrics, false otherwise. + */ + boolean isLeader(); + + /** + * Get the node ID. + * @return The node ID. + */ + int nodeId(); + + /** + * Get the object storage instance. + * @return The object storage instance. + */ + ObjectStorage objectStorage(); + + /** + * Get the base labels to include in all metrics. + * @return The base labels. + */ + List> baseLabels(); + + /** + * Get the interval in milliseconds for metrics export. + * @return The interval in milliseconds. + */ + int intervalMs(); +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporter.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporter.java new file mode 100644 index 0000000000..83d6d24b8f --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporter.java @@ -0,0 +1,29 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter; + +import io.opentelemetry.sdk.metrics.export.MetricReader; + +/** + * An interface for metrics exporters, which can be converted to an OpenTelemetry MetricReader. + */ +public interface MetricsExporter { + MetricReader asMetricReader(); +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterProvider.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterProvider.java new file mode 100644 index 0000000000..e9b1771031 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterProvider.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter; + +import java.net.URI; +import java.util.List; +import java.util.Map; + +/** + * Service Provider Interface that allows extending the available metrics exporters + * without modifying the core AutoMQ OpenTelemetry module. + */ +public interface MetricsExporterProvider { + + /** + * @param scheme exporter scheme (e.g. "rw") + * @return true if this provider can create an exporter for the supplied scheme + */ + boolean supports(String scheme); + + /** + * Creates a metrics exporter for the provided URI. + * + * @param config metrics configuration + * @param uri original exporter URI + * @param queryParameters parsed query parameters from the URI + * @return a MetricsExporter instance, or {@code null} if unable to create one + */ + MetricsExporter create(MetricsExportConfig config, URI uri, Map> queryParameters); +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterType.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterType.java new file mode 100644 index 0000000000..d8cc6f97fb --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterType.java @@ -0,0 +1,46 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter; + +public enum MetricsExporterType { + OTLP("otlp"), + PROMETHEUS("prometheus"), + OPS("ops"), + OTHER("other"); + + private final String type; + + MetricsExporterType(String type) { + this.type = type; + } + + public String getType() { + return type; + } + + public static MetricsExporterType fromString(String type) { + for (MetricsExporterType exporterType : MetricsExporterType.values()) { + if (exporterType.getType().equalsIgnoreCase(type)) { + return exporterType; + } + } + return OTHER; + } +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterURI.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterURI.java new file mode 100644 index 0000000000..f7bef935a3 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/MetricsExporterURI.java @@ -0,0 +1,220 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter; + +import com.automq.opentelemetry.common.OTLPCompressionType; +import com.automq.opentelemetry.common.OTLPProtocol; + +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.URI; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.ServiceLoader; + +/** + * Parses the exporter URI and creates the corresponding MetricsExporter instances. + */ +public class MetricsExporterURI { + private static final Logger LOGGER = LoggerFactory.getLogger(MetricsExporterURI.class); + + private static final List PROVIDERS; + + static { + List providers = new ArrayList<>(); + ServiceLoader.load(MetricsExporterProvider.class).forEach(providers::add); + PROVIDERS = Collections.unmodifiableList(providers); + if (!PROVIDERS.isEmpty()) { + LOGGER.info("Loaded {} telemetry exporter providers", PROVIDERS.size()); + } + } + + private final List metricsExporters; + + private MetricsExporterURI(List metricsExporters) { + this.metricsExporters = metricsExporters != null ? metricsExporters : new ArrayList<>(); + } + + public List getMetricsExporters() { + return metricsExporters; + } + + public static MetricsExporterURI parse(String uriStr, MetricsExportConfig config) { + LOGGER.info("Parsing metrics exporter URI: {}", uriStr); + if (StringUtils.isBlank(uriStr)) { + LOGGER.info("Metrics exporter URI is not configured, no metrics will be exported."); + return new MetricsExporterURI(Collections.emptyList()); + } + + // Support multiple exporters separated by comma + String[] exporterUris = uriStr.split(","); + if (exporterUris.length == 0) { + return new MetricsExporterURI(Collections.emptyList()); + } + + List exporters = new ArrayList<>(); + for (String uri : exporterUris) { + if (StringUtils.isBlank(uri)) { + continue; + } + MetricsExporter exporter = parseExporter(config, uri.trim()); + if (exporter != null) { + exporters.add(exporter); + } + } + return new MetricsExporterURI(exporters); + } + + public static MetricsExporter parseExporter(MetricsExportConfig config, String uriStr) { + try { + URI uri = new URI(uriStr); + String type = uri.getScheme(); + if (StringUtils.isBlank(type)) { + LOGGER.error("Invalid metrics exporter URI: {}, exporter scheme is missing", uriStr); + throw new IllegalArgumentException("Invalid metrics exporter URI: " + uriStr); + } + + Map> queries = parseQueryParameters(uri); + return parseExporter(config, type, queries, uri); + } catch (Exception e) { + LOGGER.warn("Parse metrics exporter URI {} failed", uriStr, e); + throw new IllegalArgumentException("Invalid metrics exporter URI: " + uriStr, e); + } + } + + public static MetricsExporter parseExporter(MetricsExportConfig config, String type, Map> queries, URI uri) { + MetricsExporterType exporterType = MetricsExporterType.fromString(type); + switch (exporterType) { + case PROMETHEUS: + return buildPrometheusExporter(config, queries, uri); + case OTLP: + return buildOtlpExporter(config, queries, uri); + case OPS: + return buildS3MetricsExporter(config, uri); + default: + break; + } + + MetricsExporterProvider provider = findProvider(type); + if (provider != null) { + MetricsExporter exporter = provider.create(config, uri, queries); + if (exporter != null) { + return exporter; + } + } + + LOGGER.warn("Unsupported metrics exporter type: {}", type); + return null; + } + + private static MetricsExporter buildPrometheusExporter(MetricsExportConfig config, Map> queries, URI uri) { + // Use query parameters if available, otherwise fall back to URI authority or config defaults + String host = getStringFromQuery(queries, "host", uri.getHost()); + if (StringUtils.isBlank(host)) { + host = "localhost"; + } + + int port = uri.getPort(); + if (port <= 0) { + String portStr = getStringFromQuery(queries, "port", null); + if (StringUtils.isNotBlank(portStr)) { + try { + port = Integer.parseInt(portStr); + } catch (NumberFormatException e) { + LOGGER.warn("Invalid port in query parameters: {}, using default", portStr); + port = 9090; + } + } else { + port = 9090; + } + } + + return new PrometheusMetricsExporter(host, port, config.baseLabels()); + } + + private static MetricsExporter buildOtlpExporter(MetricsExportConfig config, Map> queries, URI uri) { + // Get endpoint from query parameters or construct from URI + String endpoint = getStringFromQuery(queries, "endpoint", null); + if (StringUtils.isBlank(endpoint)) { + endpoint = uri.getScheme() + "://" + uri.getAuthority(); + } + + // Get protocol from query parameters or config + String protocol = getStringFromQuery(queries, "protocol", OTLPProtocol.GRPC.getProtocol()); + + // Get compression from query parameters or config + String compression = getStringFromQuery(queries, "compression", OTLPCompressionType.NONE.getType()); + + return new OTLPMetricsExporter(config.intervalMs(), endpoint, protocol, compression); + } + + private static MetricsExporter buildS3MetricsExporter(MetricsExportConfig config, URI uri) { + LOGGER.info("Creating S3 metrics exporter from URI: {}", uri); + if (config.objectStorage() == null) { + LOGGER.warn("No object storage configured, skip s3 metrics exporter creation."); + return null; + } + // Create the S3MetricsExporterAdapter with appropriate configuration + return new com.automq.opentelemetry.exporter.s3.S3MetricsExporterAdapter(config); + } + + private static Map> parseQueryParameters(URI uri) { + Map> queries = new HashMap<>(); + String query = uri.getQuery(); + if (StringUtils.isNotBlank(query)) { + String[] pairs = query.split("&"); + for (String pair : pairs) { + String[] keyValue = pair.split("=", 2); + if (keyValue.length == 2) { + String key = keyValue[0]; + String value = keyValue[1]; + queries.computeIfAbsent(key, k -> new ArrayList<>()).add(value); + } + } + } + return queries; + } + + private static String getStringFromQuery(Map> queries, String key, String defaultValue) { + List values = queries.get(key); + if (values != null && !values.isEmpty()) { + return values.get(0); + } + return defaultValue; + } + + private static MetricsExporterProvider findProvider(String scheme) { + for (MetricsExporterProvider provider : PROVIDERS) { + try { + if (provider.supports(scheme)) { + return provider; + } + } catch (Exception e) { + LOGGER.warn("Telemetry exporter provider {} failed to evaluate support for scheme {}", provider.getClass().getName(), scheme, e); + } + } + return null; + } +} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPMetricsExporter.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/OTLPMetricsExporter.java similarity index 51% rename from core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPMetricsExporter.java rename to automq-metrics/src/main/java/com/automq/opentelemetry/exporter/OTLPMetricsExporter.java index 3c17308591..063838009c 100644 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPMetricsExporter.java +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/OTLPMetricsExporter.java @@ -1,18 +1,28 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package kafka.log.stream.s3.telemetry.exporter; +package com.automq.opentelemetry.exporter; -import org.apache.kafka.common.utils.Utils; +import com.automq.opentelemetry.common.OTLPCompressionType; +import com.automq.opentelemetry.common.OTLPProtocol; +import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -28,13 +38,16 @@ public class OTLPMetricsExporter implements MetricsExporter { private static final Logger LOGGER = LoggerFactory.getLogger(OTLPMetricsExporter.class); - private final int intervalMs; + private final long intervalMs; private final String endpoint; private final OTLPProtocol protocol; private final OTLPCompressionType compression; + // Default timeout for OTLP exporters + private static final long DEFAULT_EXPORTER_TIMEOUT_MS = 30000; + - public OTLPMetricsExporter(int intervalMs, String endpoint, String protocol, String compression) { - if (Utils.isBlank(endpoint) || "null".equals(endpoint)) { + public OTLPMetricsExporter(long intervalMs, String endpoint, String protocol, String compression) { + if (StringUtils.isBlank(endpoint) || "null".equals(endpoint)) { throw new IllegalArgumentException("OTLP endpoint is required"); } this.intervalMs = intervalMs; @@ -42,7 +55,7 @@ public OTLPMetricsExporter(int intervalMs, String endpoint, String protocol, Str this.protocol = OTLPProtocol.fromString(protocol); this.compression = OTLPCompressionType.fromString(compression); LOGGER.info("OTLPMetricsExporter initialized with endpoint: {}, protocol: {}, compression: {}, intervalMs: {}", - endpoint, protocol, compression, intervalMs); + endpoint, protocol, compression, intervalMs); } public String endpoint() { @@ -57,31 +70,29 @@ public OTLPCompressionType compression() { return compression; } - public int intervalMs() { + public long intervalMs() { return intervalMs; } @Override public MetricReader asMetricReader() { - PeriodicMetricReaderBuilder builder; - switch (protocol) { - case GRPC: + PeriodicMetricReaderBuilder builder = switch (protocol) { + case GRPC -> { OtlpGrpcMetricExporterBuilder otlpExporterBuilder = OtlpGrpcMetricExporter.builder() .setEndpoint(endpoint) .setCompression(compression.getType()) - .setTimeout(Duration.ofMillis(ExporterConstants.DEFAULT_EXPORTER_TIMEOUT_MS)); - builder = PeriodicMetricReader.builder(otlpExporterBuilder.build()); - break; - case HTTP: + .setTimeout(Duration.ofMillis(DEFAULT_EXPORTER_TIMEOUT_MS)); + yield PeriodicMetricReader.builder(otlpExporterBuilder.build()); + } + case HTTP -> { OtlpHttpMetricExporterBuilder otlpHttpExporterBuilder = OtlpHttpMetricExporter.builder() .setEndpoint(endpoint) .setCompression(compression.getType()) - .setTimeout(Duration.ofMillis(ExporterConstants.DEFAULT_EXPORTER_TIMEOUT_MS)); - builder = PeriodicMetricReader.builder(otlpHttpExporterBuilder.build()); - break; - default: - throw new IllegalArgumentException("Unsupported OTLP protocol: " + protocol); - } + .setTimeout(Duration.ofMillis(DEFAULT_EXPORTER_TIMEOUT_MS)); + yield PeriodicMetricReader.builder(otlpHttpExporterBuilder.build()); + } + default -> throw new IllegalArgumentException("Unsupported OTLP protocol: " + protocol); + }; return builder.setInterval(Duration.ofMillis(intervalMs)).build(); } diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/PrometheusMetricsExporter.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/PrometheusMetricsExporter.java new file mode 100644 index 0000000000..f428870fe7 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/PrometheusMetricsExporter.java @@ -0,0 +1,68 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter; + +import com.automq.opentelemetry.TelemetryConstants; + +import org.apache.commons.lang3.tuple.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +import io.opentelemetry.exporter.prometheus.PrometheusHttpServer; +import io.opentelemetry.sdk.metrics.export.MetricReader; + +public class PrometheusMetricsExporter implements MetricsExporter { + private static final Logger LOGGER = LoggerFactory.getLogger(PrometheusMetricsExporter.class); + private final String host; + private final int port; + private final Set baseLabelKeys; + + public PrometheusMetricsExporter(String host, int port, List> baseLabels) { + if (host == null || host.isEmpty()) { + throw new IllegalArgumentException("Illegal Prometheus host"); + } + if (port <= 0) { + throw new IllegalArgumentException("Illegal Prometheus port"); + } + this.host = host; + this.port = port; + this.baseLabelKeys = baseLabels.stream().map(Pair::getKey).collect(Collectors.toSet()); + LOGGER.info("PrometheusMetricsExporter initialized with host: {}, port: {}, labels: {}", host, port, baseLabels); + } + + @Override + public MetricReader asMetricReader() { + return PrometheusHttpServer.builder() + .setHost(host) + .setPort(port) + // This filter is to align with the original behavior, allowing only specific resource attributes + // to be converted to prometheus labels. + .setAllowedResourceAttributesFilter(resourceAttributeKey -> + TelemetryConstants.PROMETHEUS_JOB_KEY.equals(resourceAttributeKey) + || TelemetryConstants.PROMETHEUS_INSTANCE_KEY.equals(resourceAttributeKey) + || TelemetryConstants.HOST_NAME_KEY.equals(resourceAttributeKey) + || baseLabelKeys.contains(resourceAttributeKey)) + .build(); + } +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/CompressionUtils.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/CompressionUtils.java new file mode 100644 index 0000000000..20afdd6b36 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/CompressionUtils.java @@ -0,0 +1,86 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter.s3; + +import com.automq.stream.s3.ByteBufAlloc; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +import io.netty.buffer.ByteBuf; + +/** + * Utility class for data compression and decompression. + */ +public class CompressionUtils { + + /** + * Compress a ByteBuf using GZIP. + * + * @param input The input ByteBuf to compress. + * @return A new ByteBuf containing the compressed data. + * @throws IOException If an I/O error occurs during compression. + */ + public static ByteBuf compress(ByteBuf input) throws IOException { + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream); + + byte[] buffer = new byte[input.readableBytes()]; + input.readBytes(buffer); + gzipOutputStream.write(buffer); + gzipOutputStream.close(); + + ByteBuf compressed = ByteBufAlloc.byteBuffer(byteArrayOutputStream.size()); + compressed.writeBytes(byteArrayOutputStream.toByteArray()); + return compressed; + } + + /** + * Decompress a GZIP-compressed ByteBuf. + * + * @param input The compressed ByteBuf to decompress. + * @return A new ByteBuf containing the decompressed data. + * @throws IOException If an I/O error occurs during decompression. + */ + public static ByteBuf decompress(ByteBuf input) throws IOException { + byte[] compressedData = new byte[input.readableBytes()]; + input.readBytes(compressedData); + ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(compressedData); + GZIPInputStream gzipInputStream = new GZIPInputStream(byteArrayInputStream); + + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = gzipInputStream.read(buffer)) != -1) { + byteArrayOutputStream.write(buffer, 0, bytesRead); + } + + gzipInputStream.close(); + byteArrayOutputStream.close(); + + byte[] uncompressedData = byteArrayOutputStream.toByteArray(); + ByteBuf output = ByteBufAlloc.byteBuffer(uncompressedData.length); + output.writeBytes(uncompressedData); + return output; + } +} diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/PrometheusUtils.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/PrometheusUtils.java new file mode 100644 index 0000000000..3d4a671a08 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/PrometheusUtils.java @@ -0,0 +1,276 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter.s3; + +import org.apache.commons.lang3.StringUtils; + +import java.util.Locale; + +/** + * Utility class for Prometheus metric and label naming. + */ +public class PrometheusUtils { + private static final String TOTAL_SUFFIX = "_total"; + + /** + * Get the Prometheus unit from the OpenTelemetry unit. + * + * @param unit The OpenTelemetry unit. + * @return The Prometheus unit. + */ + public static String getPrometheusUnit(String unit) { + if (unit.contains("{")) { + return ""; + } + switch (unit) { + // Time + case "d": + return "days"; + case "h": + return "hours"; + case "min": + return "minutes"; + case "s": + return "seconds"; + case "ms": + return "milliseconds"; + case "us": + return "microseconds"; + case "ns": + return "nanoseconds"; + // Bytes + case "By": + return "bytes"; + case "KiBy": + return "kibibytes"; + case "MiBy": + return "mebibytes"; + case "GiBy": + return "gibibytes"; + case "TiBy": + return "tibibytes"; + case "KBy": + return "kilobytes"; + case "MBy": + return "megabytes"; + case "GBy": + return "gigabytes"; + case "TBy": + return "terabytes"; + // SI + case "m": + return "meters"; + case "V": + return "volts"; + case "A": + return "amperes"; + case "J": + return "joules"; + case "W": + return "watts"; + case "g": + return "grams"; + // Misc + case "Cel": + return "celsius"; + case "Hz": + return "hertz"; + case "1": + return ""; + case "%": + return "percent"; + // Rate units (per second) + case "1/s": + return "per_second"; + case "By/s": + return "bytes_per_second"; + case "KiBy/s": + return "kibibytes_per_second"; + case "MiBy/s": + return "mebibytes_per_second"; + case "GiBy/s": + return "gibibytes_per_second"; + case "KBy/s": + return "kilobytes_per_second"; + case "MBy/s": + return "megabytes_per_second"; + case "GBy/s": + return "gigabytes_per_second"; + // Rate units (per minute) + case "1/min": + return "per_minute"; + case "By/min": + return "bytes_per_minute"; + // Rate units (per hour) + case "1/h": + return "per_hour"; + case "By/h": + return "bytes_per_hour"; + // Rate units (per day) + case "1/d": + return "per_day"; + case "By/d": + return "bytes_per_day"; + default: + return unit; + } + } + + /** + * Map a metric name to a Prometheus-compatible name. + * + * @param name The original metric name. + * @param unit The metric unit. + * @param isCounter Whether the metric is a counter. + * @param isGauge Whether the metric is a gauge. + * @return The Prometheus-compatible metric name. + */ + public static String mapMetricsName(String name, String unit, boolean isCounter, boolean isGauge) { + // Replace "." into "_" + name = name.replaceAll("\\.", "_"); + + String prometheusUnit = getPrometheusUnit(unit); + boolean shouldAppendUnit = StringUtils.isNotBlank(prometheusUnit) && !name.contains(prometheusUnit); + + // append prometheus unit if not null or empty. + // unit should be appended before type suffix + if (shouldAppendUnit) { + name = name + "_" + prometheusUnit; + } + + // trim counter's _total suffix so the unit is placed before it. + if (isCounter && name.endsWith(TOTAL_SUFFIX)) { + name = name.substring(0, name.length() - TOTAL_SUFFIX.length()); + } + + // replace _total suffix, or add if it wasn't already present. + if (isCounter) { + name = name + TOTAL_SUFFIX; + } + + // special case - gauge with intelligent Connect metric handling + if ("1".equals(unit) && isGauge && !name.contains("ratio")) { + if (isConnectMetric(name)) { + // For Connect metrics, use improved logic to avoid misleading _ratio suffix + if (shouldAddRatioSuffixForConnect(name)) { + name = name + "_ratio"; + } + } else { + // For other metrics, maintain original behavior + name = name + "_ratio"; + } + } + return name; + } + + /** + * Map a label name to a Prometheus-compatible name. + * + * @param name The original label name. + * @return The Prometheus-compatible label name. + */ + public static String mapLabelName(String name) { + if (StringUtils.isBlank(name)) { + return ""; + } + return name.replaceAll("\\.", "_"); + } + + /** + * Check if a metric name is related to Kafka Connect. + * + * @param name The metric name to check. + * @return true if it's a Connect metric, false otherwise. + */ + private static boolean isConnectMetric(String name) { + String lowerName = name.toLowerCase(Locale.ROOT); + return lowerName.contains("kafka_connector_") || + lowerName.contains("kafka_task_") || + lowerName.contains("kafka_worker_") || + lowerName.contains("kafka_connect_") || + lowerName.contains("kafka_source_task_") || + lowerName.contains("kafka_sink_task_") || + lowerName.contains("connector_metrics") || + lowerName.contains("task_metrics") || + lowerName.contains("worker_metrics") || + lowerName.contains("source_task_metrics") || + lowerName.contains("sink_task_metrics"); + } + + /** + * Intelligently determine if a Connect metric should have a _ratio suffix. + * This method avoids adding misleading _ratio suffixes to count-based metrics. + * + * @param name The metric name to check. + * @return true if _ratio suffix should be added, false otherwise. + */ + private static boolean shouldAddRatioSuffixForConnect(String name) { + String lowerName = name.toLowerCase(Locale.ROOT); + + if (hasRatioRelatedWords(lowerName)) { + return false; + } + + if (isCountMetric(lowerName)) { + return false; + } + + return isRatioMetric(lowerName); + } + + private static boolean hasRatioRelatedWords(String lowerName) { + return lowerName.contains("ratio") || lowerName.contains("percent") || + lowerName.contains("rate") || lowerName.contains("fraction"); + } + + private static boolean isCountMetric(String lowerName) { + return hasBasicCountKeywords(lowerName) || hasConnectCountKeywords(lowerName) || + hasStatusCountKeywords(lowerName); + } + + private static boolean hasBasicCountKeywords(String lowerName) { + return lowerName.contains("count") || lowerName.contains("num") || + lowerName.contains("size") || lowerName.contains("total") || + lowerName.contains("active") || lowerName.contains("current"); + } + + private static boolean hasConnectCountKeywords(String lowerName) { + return lowerName.contains("partition") || lowerName.contains("task") || + lowerName.contains("connector") || lowerName.contains("seq_no") || + lowerName.contains("seq_num") || lowerName.contains("attempts"); + } + + private static boolean hasStatusCountKeywords(String lowerName) { + return lowerName.contains("success") || lowerName.contains("failure") || + lowerName.contains("errors") || lowerName.contains("retries") || + lowerName.contains("skipped") || lowerName.contains("running") || + lowerName.contains("paused") || lowerName.contains("failed") || + lowerName.contains("destroyed"); + } + + private static boolean isRatioMetric(String lowerName) { + return lowerName.contains("utilization") || + lowerName.contains("usage") || + lowerName.contains("load") || + lowerName.contains("efficiency") || + lowerName.contains("hit_rate") || + lowerName.contains("miss_rate"); + } +} diff --git a/automq-shell/src/main/java/com/automq/shell/metrics/S3MetricsExporter.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/S3MetricsExporter.java similarity index 88% rename from automq-shell/src/main/java/com/automq/shell/metrics/S3MetricsExporter.java rename to automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/S3MetricsExporter.java index 7d501dfaae..432abeea28 100644 --- a/automq-shell/src/main/java/com/automq/shell/metrics/S3MetricsExporter.java +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/S3MetricsExporter.java @@ -1,20 +1,30 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package com.automq.shell.metrics; +package com.automq.opentelemetry.exporter.s3; +import com.automq.opentelemetry.exporter.MetricsExportConfig; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorage.ObjectInfo; import com.automq.stream.s3.operator.ObjectStorage.ObjectPath; import com.automq.stream.s3.operator.ObjectStorage.WriteOptions; +import com.automq.stream.utils.Threads; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -50,6 +60,9 @@ import io.opentelemetry.sdk.metrics.data.MetricData; import io.opentelemetry.sdk.metrics.export.MetricExporter; +/** + * An S3 metrics exporter that uploads metrics data to S3 buckets. + */ public class S3MetricsExporter implements MetricExporter { private static final Logger LOGGER = LoggerFactory.getLogger(S3MetricsExporter.class); @@ -58,13 +71,13 @@ public class S3MetricsExporter implements MetricExporter { public static final int MAX_JITTER_INTERVAL = 60 * 1000; public static final int DEFAULT_BUFFER_SIZE = 16 * 1024 * 1024; - private final S3MetricsConfig config; + private final MetricsExportConfig config; private final Map defaultTagMap = new HashMap<>(); private final ByteBuf uploadBuffer = Unpooled.directBuffer(DEFAULT_BUFFER_SIZE); - private final Random random = new Random(); + private static final Random RANDOM = new Random(); private volatile long lastUploadTimestamp = System.currentTimeMillis(); - private volatile long nextUploadInterval = UPLOAD_INTERVAL + random.nextInt(MAX_JITTER_INTERVAL); + private volatile long nextUploadInterval = UPLOAD_INTERVAL + RANDOM.nextInt(MAX_JITTER_INTERVAL); private final ObjectStorage objectStorage; private final ObjectMapper objectMapper = new ObjectMapper(); @@ -73,7 +86,12 @@ public class S3MetricsExporter implements MetricExporter { private final Thread uploadThread; private final Thread cleanupThread; - public S3MetricsExporter(S3MetricsConfig config) { + /** + * Creates a new S3MetricsExporter. + * + * @param config The configuration for the S3 metrics exporter. + */ + public S3MetricsExporter(MetricsExportConfig config) { this.config = config; this.objectStorage = config.objectStorage(); @@ -91,6 +109,9 @@ public S3MetricsExporter(S3MetricsConfig config) { cleanupThread.setDaemon(true); } + /** + * Starts the exporter threads. + */ public void start() { uploadThread.start(); cleanupThread.start(); @@ -129,7 +150,7 @@ private class CleanupTask implements Runnable { public void run() { while (!Thread.currentThread().isInterrupted()) { try { - if (closed || !config.isActiveController()) { + if (closed || !config.isLeader()) { Thread.sleep(Duration.ofMinutes(1).toMillis()); continue; } @@ -152,8 +173,7 @@ public void run() { CompletableFuture.allOf(deleteFutures).join(); } } - - Thread.sleep(Duration.ofMinutes(1).toMillis()); + Threads.sleep(Duration.ofMinutes(1).toMillis()); } catch (InterruptedException e) { break; } catch (Exception e) { @@ -242,13 +262,13 @@ public CompletableResultCode flush() { synchronized (uploadBuffer) { if (uploadBuffer.readableBytes() > 0) { try { - objectStorage.write(WriteOptions.DEFAULT, getObjectKey(), uploadBuffer.retainedSlice().asReadOnly()).get(); + objectStorage.write(WriteOptions.DEFAULT, getObjectKey(), CompressionUtils.compress(uploadBuffer.slice().asReadOnly())).get(); } catch (Exception e) { LOGGER.error("Failed to upload metrics to s3", e); return CompletableResultCode.ofFailure(); } finally { lastUploadTimestamp = System.currentTimeMillis(); - nextUploadInterval = UPLOAD_INTERVAL + random.nextInt(MAX_JITTER_INTERVAL); + nextUploadInterval = UPLOAD_INTERVAL + RANDOM.nextInt(MAX_JITTER_INTERVAL); uploadBuffer.clear(); } } diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/S3MetricsExporterAdapter.java b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/S3MetricsExporterAdapter.java new file mode 100644 index 0000000000..b7f916cda3 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/exporter/s3/S3MetricsExporterAdapter.java @@ -0,0 +1,63 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.exporter.s3; + +import com.automq.opentelemetry.exporter.MetricsExportConfig; +import com.automq.opentelemetry.exporter.MetricsExporter; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; + +import io.opentelemetry.sdk.metrics.export.MetricReader; +import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader; + +/** + * An adapter class that implements the MetricsExporter interface and uses S3MetricsExporter + * for actual metrics exporting functionality. + */ +public class S3MetricsExporterAdapter implements MetricsExporter { + private static final Logger LOGGER = LoggerFactory.getLogger(S3MetricsExporterAdapter.class); + + private final MetricsExportConfig metricsExportConfig; + + /** + * Creates a new S3MetricsExporterAdapter. + * + * @param metricsExportConfig The configuration for the S3 metrics exporter. + */ + public S3MetricsExporterAdapter(MetricsExportConfig metricsExportConfig) { + this.metricsExportConfig = metricsExportConfig; + LOGGER.info("S3MetricsExporterAdapter initialized with labels :{}", metricsExportConfig.baseLabels()); + } + + @Override + public MetricReader asMetricReader() { + // Create and start the S3MetricsExporter + S3MetricsExporter s3MetricsExporter = new S3MetricsExporter(metricsExportConfig); + s3MetricsExporter.start(); + + // Create and return the periodic metric reader + return PeriodicMetricReader.builder(s3MetricsExporter) + .setInterval(Duration.ofMillis(metricsExportConfig.intervalMs())) + .build(); + } +} \ No newline at end of file diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/DeltaHistogram.java b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/DeltaHistogram.java similarity index 72% rename from core/src/main/scala/kafka/log/stream/s3/telemetry/otel/DeltaHistogram.java rename to automq-metrics/src/main/java/com/automq/opentelemetry/yammer/DeltaHistogram.java index 6244baa37a..8f4fd459f5 100644 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/DeltaHistogram.java +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/DeltaHistogram.java @@ -1,15 +1,23 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package kafka.log.stream.s3.telemetry.otel; +package com.automq.opentelemetry.yammer; import com.yammer.metrics.core.Histogram; import com.yammer.metrics.core.Timer; diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelMetricUtils.java b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/OTelMetricUtils.java similarity index 88% rename from core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelMetricUtils.java rename to automq-metrics/src/main/java/com/automq/opentelemetry/yammer/OTelMetricUtils.java index bed205156f..7d58de2661 100644 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelMetricUtils.java +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/OTelMetricUtils.java @@ -1,15 +1,23 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package kafka.log.stream.s3.telemetry.otel; +package com.automq.opentelemetry.yammer; import com.yammer.metrics.core.MetricName; diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelMetricsProcessor.java b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/YammerMetricsProcessor.java similarity index 51% rename from core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelMetricsProcessor.java rename to automq-metrics/src/main/java/com/automq/opentelemetry/yammer/YammerMetricsProcessor.java index dde58d44b9..0875ccae2f 100644 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelMetricsProcessor.java +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/YammerMetricsProcessor.java @@ -1,17 +1,24 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package kafka.log.stream.s3.telemetry.otel; +package com.automq.opentelemetry.yammer; -import kafka.autobalancer.metricsreporter.metric.MetricsUtils; import com.yammer.metrics.core.Counter; import com.yammer.metrics.core.Gauge; @@ -24,16 +31,54 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collections; +import java.util.HashMap; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.common.AttributesBuilder; import io.opentelemetry.api.metrics.Meter; -import scala.UninitializedFieldError; -public class OTelMetricsProcessor implements MetricProcessor { - private static final Logger LOGGER = LoggerFactory.getLogger(OTelMetricsProcessor.class); +/** + * A metrics processor that bridges Yammer metrics to OpenTelemetry metrics. + * + *

This processor specifically handles Histogram and Timer metrics from the Yammer metrics + * library and converts them to OpenTelemetry gauge metrics that track delta mean values. + * It implements the Yammer {@link MetricProcessor} interface to process metrics and creates + * corresponding OpenTelemetry metrics with proper attributes derived from the metric scope. + * + *

The processor: + *

    + *
  • Converts Yammer Histogram and Timer metrics to OpenTelemetry gauges
    • + *
  • Calculates delta mean values using {@link DeltaHistogram}
    • + *
  • Parses metric scopes to extract attributes for OpenTelemetry metrics
    • + *
  • Maintains a registry of processed metrics for lifecycle management
    • + *
  • Supports metric removal when metrics are no longer needed
    • + *
+ * + *

Supported metric types: + *

    + *
  • {@link Histogram} - Converted to delta mean gauge
    • + *
  • {@link Timer} - Converted to delta mean gauge
    • + *
+ * + *

Unsupported metric types (will throw {@link UnsupportedOperationException}): + *

    + *
  • {@link Counter}
    • + *
  • {@link Gauge}
    • + *
  • {@link Metered}
    • + *
+ * + *

Thread Safety: This class is thread-safe and uses concurrent data structures + * to handle metrics registration and removal from multiple threads. + * + * @see MetricProcessor + * @see DeltaHistogram + * @see OTelMetricUtils + */ +public class YammerMetricsProcessor implements MetricProcessor { + private static final Logger LOGGER = LoggerFactory.getLogger(YammerMetricsProcessor.class); private final Map> metrics = new ConcurrentHashMap<>(); private Meter meter = null; @@ -63,9 +108,9 @@ public void processTimer(MetricName name, Timer timer, Void unused) { private void processDeltaHistogramMetric(MetricName name, DeltaHistogram deltaHistogram) { if (meter == null) { - throw new UninitializedFieldError("Meter is not initialized"); + throw new IllegalStateException("Meter is not initialized"); } - Map tags = MetricsUtils.yammerMetricScopeToTags(name.getScope()); + Map tags = yammerMetricScopeToTags(name.getScope()); AttributesBuilder attrBuilder = Attributes.builder(); if (tags != null) { String value = tags.remove(OTelMetricUtils.REQUEST_TAG_KEY); @@ -108,6 +153,29 @@ public void remove(MetricName metricName) { }); } + /** + * Convert a yammer metrics scope to a tags map. + * + * @param scope Scope of the Yammer metric. + * @return Empty map for {@code null} scope, {@code null} for scope with keys without a matching value (i.e. unacceptable + * scope) (see ...), parsed tags otherwise. + */ + public static Map yammerMetricScopeToTags(String scope) { + if (scope != null) { + String[] kv = scope.split("\\."); + if (kv.length % 2 != 0) { + return null; + } + Map tags = new HashMap<>(); + for (int i = 0; i < kv.length; i += 2) { + tags.put(kv[i], kv[i + 1]); + } + return tags; + } else { + return Collections.emptyMap(); + } + } + static class MetricWrapper { private final Attributes attr; private final DeltaHistogram deltaHistogram; diff --git a/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/YammerMetricsReporter.java b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/YammerMetricsReporter.java new file mode 100644 index 0000000000..9482a08d68 --- /dev/null +++ b/automq-metrics/src/main/java/com/automq/opentelemetry/yammer/YammerMetricsReporter.java @@ -0,0 +1,93 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.opentelemetry.yammer; + +import com.yammer.metrics.core.Metric; +import com.yammer.metrics.core.MetricName; +import com.yammer.metrics.core.MetricsRegistry; +import com.yammer.metrics.core.MetricsRegistryListener; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.Closeable; +import java.io.IOException; + +import io.opentelemetry.api.metrics.Meter; + +/** + * A listener that bridges Yammer Histogram metrics to OpenTelemetry. + * It listens for new metrics added to a MetricsRegistry and creates corresponding + * OTel gauge metrics for mean and max values of histograms. + */ +public class YammerMetricsReporter implements MetricsRegistryListener, Closeable { + private static final Logger LOGGER = LoggerFactory.getLogger(YammerMetricsReporter.class); + private final MetricsRegistry metricsRegistry; + private final YammerMetricsProcessor metricsProcessor; + private volatile Meter meter; + + public YammerMetricsReporter(MetricsRegistry metricsRegistry) { + this.metricsRegistry = metricsRegistry; + this.metricsProcessor = new YammerMetricsProcessor(); + } + + public void start(Meter meter) { + this.meter = meter; + this.metricsProcessor.init(meter); + metricsRegistry.addListener(this); + LOGGER.info("OTelHistogramReporter started"); + } + + @Override + public void onMetricAdded(MetricName name, Metric metric) { + if (OTelMetricUtils.isInterestedMetric(name)) { + if (this.meter == null) { + LOGGER.info("Not initialized yet, skipping metric: {}", name); + return; + } + try { + metric.processWith(this.metricsProcessor, name, null); + } catch (Throwable t) { + LOGGER.error("Failed to process metric: {}", name, t); + } + } + } + + @Override + public void onMetricRemoved(MetricName name) { + try { + this.metricsProcessor.remove(name); + } catch (Throwable ignored) { + + } + } + + @Override + public void close() throws IOException { + try { + // Remove this reporter as a listener from the metrics registry + metricsRegistry.removeListener(this); + LOGGER.info("YammerMetricsReporter stopped and removed from metrics registry"); + } catch (Exception e) { + LOGGER.error("Error while closing YammerMetricsReporter", e); + throw new IOException("Failed to close YammerMetricsReporter", e); + } + } +} \ No newline at end of file diff --git a/automq-shell/build.gradle b/automq-shell/build.gradle index 4e8b5d9510..735758f5ce 100644 --- a/automq-shell/build.gradle +++ b/automq-shell/build.gradle @@ -18,7 +18,8 @@ dependencies { compileOnly libs.awsSdkAuth implementation libs.reload4j implementation libs.nettyBuffer - implementation libs.opentelemetrySdk + implementation project(':automq-metrics') + implementation project(':automq-log-uploader') implementation libs.jacksonDatabind implementation libs.jacksonYaml implementation libs.commonLang @@ -65,4 +66,4 @@ jar { manifest { attributes 'Main-Class': 'com.automq.shell.AutoMQCLI' } -} \ No newline at end of file +} diff --git a/automq-shell/src/main/java/com/automq/shell/AutoMQApplication.java b/automq-shell/src/main/java/com/automq/shell/AutoMQApplication.java index 06620aa296..557ce79ce1 100644 --- a/automq-shell/src/main/java/com/automq/shell/AutoMQApplication.java +++ b/automq-shell/src/main/java/com/automq/shell/AutoMQApplication.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell; @@ -33,9 +41,8 @@ public static boolean registerSingleton(Class type, T singleton, boolean if (override) { CONTAINER.put(type, singleton); return true; - } else { - return CONTAINER.putIfAbsent(type, singleton) == null; } + return CONTAINER.putIfAbsent(type, singleton) == null; } public static T getBean(Class type) { diff --git a/automq-shell/src/main/java/com/automq/shell/AutoMQCLI.java b/automq-shell/src/main/java/com/automq/shell/AutoMQCLI.java index 5c703de10a..29fde1625e 100644 --- a/automq-shell/src/main/java/com/automq/shell/AutoMQCLI.java +++ b/automq-shell/src/main/java/com/automq/shell/AutoMQCLI.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell; diff --git a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Cluster.java b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Cluster.java index a8064221a9..1aed8ca421 100644 --- a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Cluster.java +++ b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Cluster.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.commands.cluster; diff --git a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Create.java b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Create.java index dfe89603cb..e15453c5cc 100644 --- a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Create.java +++ b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Create.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.commands.cluster; diff --git a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Deploy.java b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Deploy.java index b552b509ac..ed23fdddbc 100644 --- a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Deploy.java +++ b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Deploy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.commands.cluster; @@ -102,9 +110,11 @@ private static void bucketReadinessCheck(ClusterTopology topo) { String globalAccessKey = null; String globalSecretKey = null; for (Env env : topo.getGlobal().getEnvs()) { - if ("KAFKA_S3_ACCESS_KEY".equals(env.getName())) { + if ("KAFKA_S3_ACCESS_KEY".equals(env.getName()) || + "AWS_ACCESS_KEY_ID".equals(env.getName())) { globalAccessKey = env.getValue(); - } else if ("KAFKA_S3_SECRET_KEY".equals(env.getName())) { + } else if ("KAFKA_S3_SECRET_KEY".equals(env.getName()) || + "AWS_SECRET_ACCESS_KEY".equals(env.getName())) { globalSecretKey = env.getValue(); } } @@ -159,6 +169,7 @@ private static void appendCommonConfigsOverride(StringBuilder sb, ClusterTopolog sb.append("--override cluster.id=").append(topo.getGlobal().getClusterId()).append(" "); sb.append("--override node.id=").append(node.getNodeId()).append(" "); sb.append("--override controller.quorum.voters=").append(getQuorumVoters(topo)).append(" "); + sb.append("--override controller.quorum.bootstrap.servers=").append(getBootstrapServers(topo)).append(" "); sb.append("--override advertised.listeners=").append("PLAINTEXT://").append(node.getHost()).append(":9092").append(" "); } @@ -181,4 +192,14 @@ private static String getQuorumVoters(ClusterTopology topo) { .map(node -> node.getNodeId() + "@" + node.getHost() + ":9093") .collect(Collectors.joining(",")); } + + private static String getBootstrapServers(ClusterTopology topo) { + List nodes = topo.getControllers(); + if (!(nodes.size() == 1 || nodes.size() == 3)) { + throw new IllegalArgumentException("Only support 1 or 3 controllers"); + } + return nodes.stream() + .map(node -> node.getHost() + ":9093") + .collect(Collectors.joining(",")); + } } diff --git a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Describe.java b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Describe.java index 2064403a2a..73f85b18bc 100644 --- a/automq-shell/src/main/java/com/automq/shell/commands/cluster/Describe.java +++ b/automq-shell/src/main/java/com/automq/shell/commands/cluster/Describe.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.commands.cluster; diff --git a/automq-shell/src/main/java/com/automq/shell/constant/ServerConfigKey.java b/automq-shell/src/main/java/com/automq/shell/constant/ServerConfigKey.java index 50b0f28a5a..8da4d7b182 100644 --- a/automq-shell/src/main/java/com/automq/shell/constant/ServerConfigKey.java +++ b/automq-shell/src/main/java/com/automq/shell/constant/ServerConfigKey.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.constant; diff --git a/automq-shell/src/main/java/com/automq/shell/log/S3LogConfig.java b/automq-shell/src/main/java/com/automq/shell/log/S3LogConfig.java deleted file mode 100644 index 2842fe0a99..0000000000 --- a/automq-shell/src/main/java/com/automq/shell/log/S3LogConfig.java +++ /dev/null @@ -1,27 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.shell.log; - -import com.automq.stream.s3.operator.ObjectStorage; - -public interface S3LogConfig { - - boolean isEnabled(); - - boolean isActiveController(); - - String clusterId(); - - int nodeId(); - - ObjectStorage objectStorage(); -} diff --git a/automq-shell/src/main/java/com/automq/shell/log/S3RollingFileAppender.java b/automq-shell/src/main/java/com/automq/shell/log/S3RollingFileAppender.java deleted file mode 100644 index 04ec75daff..0000000000 --- a/automq-shell/src/main/java/com/automq/shell/log/S3RollingFileAppender.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.shell.log; - -import org.apache.log4j.RollingFileAppender; -import org.apache.log4j.spi.LoggingEvent; - -public class S3RollingFileAppender extends RollingFileAppender { - private final LogUploader logUploader = LogUploader.getInstance(); - - @Override - protected void subAppend(LoggingEvent event) { - super.subAppend(event); - if (!closed) { - LogRecorder.LogEvent logEvent = new LogRecorder.LogEvent( - event.getTimeStamp(), - event.getLevel().toString(), - event.getLoggerName(), - event.getRenderedMessage(), - event.getThrowableStrRep()); - - try { - logEvent.validate(); - } catch (IllegalArgumentException e) { - // Drop invalid log event - errorHandler.error("Failed to validate log event", e, 0); - return; - } - - logUploader.append(logEvent); - } - } -} diff --git a/automq-shell/src/main/java/com/automq/shell/metrics/PrometheusUtils.java b/automq-shell/src/main/java/com/automq/shell/metrics/PrometheusUtils.java deleted file mode 100644 index d876dc66c9..0000000000 --- a/automq-shell/src/main/java/com/automq/shell/metrics/PrometheusUtils.java +++ /dev/null @@ -1,120 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.shell.metrics; - -import org.apache.commons.lang3.StringUtils; - -public class PrometheusUtils { - private static final String TOTAL_SUFFIX = "_total"; - - public static String getPrometheusUnit(String unit) { - if (unit.contains("{")) { - return ""; - } - switch (unit) { - // Time - case "d": - return "days"; - case "h": - return "hours"; - case "min": - return "minutes"; - case "s": - return "seconds"; - case "ms": - return "milliseconds"; - case "us": - return "microseconds"; - case "ns": - return "nanoseconds"; - // Bytes - case "By": - return "bytes"; - case "KiBy": - return "kibibytes"; - case "MiBy": - return "mebibytes"; - case "GiBy": - return "gibibytes"; - case "TiBy": - return "tibibytes"; - case "KBy": - return "kilobytes"; - case "MBy": - return "megabytes"; - case "GBy": - return "gigabytes"; - case "TBy": - return "terabytes"; - // SI - case "m": - return "meters"; - case "V": - return "volts"; - case "A": - return "amperes"; - case "J": - return "joules"; - case "W": - return "watts"; - case "g": - return "grams"; - // Misc - case "Cel": - return "celsius"; - case "Hz": - return "hertz"; - case "1": - return ""; - case "%": - return "percent"; - default: - return unit; - } - } - - public static String mapMetricsName(String name, String unit, boolean isCounter, boolean isGauge) { - // Replace "." into "_" - name = name.replaceAll("\\.", "_"); - - String prometheusUnit = getPrometheusUnit(unit); - boolean shouldAppendUnit = StringUtils.isNotBlank(prometheusUnit) && !name.contains(prometheusUnit); - - // append prometheus unit if not null or empty. - // unit should be appended before type suffix - if (shouldAppendUnit) { - name = name + "_" + prometheusUnit; - } - - // trim counter's _total suffix so the unit is placed before it. - if (isCounter && name.endsWith(TOTAL_SUFFIX)) { - name = name.substring(0, name.length() - TOTAL_SUFFIX.length()); - } - - // replace _total suffix, or add if it wasn't already present. - if (isCounter) { - name = name + TOTAL_SUFFIX; - } - // special case - gauge - if (unit.equals("1") && isGauge && !name.contains("ratio")) { - name = name + "_ratio"; - } - return name; - } - - public static String mapLabelName(String name) { - if (StringUtils.isBlank(name)) { - return ""; - } - return name.replaceAll("\\.", "_"); - } -} diff --git a/automq-shell/src/main/java/com/automq/shell/metrics/S3MetricsConfig.java b/automq-shell/src/main/java/com/automq/shell/metrics/S3MetricsConfig.java deleted file mode 100644 index 377ac1d135..0000000000 --- a/automq-shell/src/main/java/com/automq/shell/metrics/S3MetricsConfig.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.shell.metrics; - -import com.automq.stream.s3.operator.ObjectStorage; - -import org.apache.commons.lang3.tuple.Pair; - -import java.util.List; - -public interface S3MetricsConfig { - - String clusterId(); - - boolean isActiveController(); - - int nodeId(); - - ObjectStorage objectStorage(); - - List> baseLabels(); -} diff --git a/automq-shell/src/main/java/com/automq/shell/model/ClusterTopology.java b/automq-shell/src/main/java/com/automq/shell/model/ClusterTopology.java index 0d6b189bb8..1257812138 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/ClusterTopology.java +++ b/automq-shell/src/main/java/com/automq/shell/model/ClusterTopology.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/model/EndpointProtocol.java b/automq-shell/src/main/java/com/automq/shell/model/EndpointProtocol.java index 915aec4424..84ff80f42d 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/EndpointProtocol.java +++ b/automq-shell/src/main/java/com/automq/shell/model/EndpointProtocol.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/model/Env.java b/automq-shell/src/main/java/com/automq/shell/model/Env.java index e47089c1c7..287afcbe21 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/Env.java +++ b/automq-shell/src/main/java/com/automq/shell/model/Env.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/model/Global.java b/automq-shell/src/main/java/com/automq/shell/model/Global.java index 7616475f3d..f391ffc0b4 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/Global.java +++ b/automq-shell/src/main/java/com/automq/shell/model/Global.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/model/Node.java b/automq-shell/src/main/java/com/automq/shell/model/Node.java index e349fa75d9..c74a2cb081 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/Node.java +++ b/automq-shell/src/main/java/com/automq/shell/model/Node.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/model/S3Url.java b/automq-shell/src/main/java/com/automq/shell/model/S3Url.java index 20ceeaa075..bc8d0167b6 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/S3Url.java +++ b/automq-shell/src/main/java/com/automq/shell/model/S3Url.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/model/StreamTags.java b/automq-shell/src/main/java/com/automq/shell/model/StreamTags.java index 874b2725a1..d80fe43d6e 100644 --- a/automq-shell/src/main/java/com/automq/shell/model/StreamTags.java +++ b/automq-shell/src/main/java/com/automq/shell/model/StreamTags.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.model; diff --git a/automq-shell/src/main/java/com/automq/shell/stream/ClientKVClient.java b/automq-shell/src/main/java/com/automq/shell/stream/ClientKVClient.java index 5f87f74a12..0e9e14175e 100644 --- a/automq-shell/src/main/java/com/automq/shell/stream/ClientKVClient.java +++ b/automq-shell/src/main/java/com/automq/shell/stream/ClientKVClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.stream; @@ -29,7 +37,6 @@ import org.apache.kafka.common.requests.s3.PutKVsRequest; import org.apache.kafka.common.utils.Time; -import com.automq.shell.metrics.S3MetricsExporter; import com.automq.stream.api.KeyValue; import org.slf4j.Logger; @@ -40,7 +47,7 @@ import java.util.Objects; public class ClientKVClient { - private static final Logger LOGGER = LoggerFactory.getLogger(S3MetricsExporter.class); + private static final Logger LOGGER = LoggerFactory.getLogger(ClientKVClient.class); private final NetworkClient networkClient; private final Node bootstrapServer; @@ -66,7 +73,9 @@ private void connect(Node node) throws IOException { public KeyValue.Value getKV(String key) throws IOException { long now = Time.SYSTEM.milliseconds(); - LOGGER.trace("[ClientKVClient]: Get KV: {}", key); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ClientKVClient]: Get KV: {}", key); + } GetKVsRequestData data = new GetKVsRequestData() .setGetKeyRequests(List.of(new GetKVsRequestData.GetKVRequest().setKey(key))); @@ -88,7 +97,9 @@ public KeyValue.Value getKV(String key) throws IOException { public KeyValue.Value putKV(String key, byte[] value) throws IOException { long now = Time.SYSTEM.milliseconds(); - LOGGER.trace("[ClientKVClient]: put KV: {}", key); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ClientKVClient]: put KV: {}", key); + } PutKVsRequestData data = new PutKVsRequestData() .setPutKVRequests(List.of(new PutKVsRequestData.PutKVRequest().setKey(key).setValue(value))); @@ -110,7 +121,10 @@ public KeyValue.Value putKV(String key, byte[] value) throws IOException { public KeyValue.Value deleteKV(String key) throws IOException { long now = Time.SYSTEM.milliseconds(); - LOGGER.trace("[ClientKVClient]: Delete KV: {}", key); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ClientKVClient]: Delete KV: {}", key); + } + DeleteKVsRequestData data = new DeleteKVsRequestData() .setDeleteKVRequests(List.of(new DeleteKVsRequestData.DeleteKVRequest().setKey(key))); diff --git a/automq-shell/src/main/java/com/automq/shell/stream/ClientStreamManager.java b/automq-shell/src/main/java/com/automq/shell/stream/ClientStreamManager.java index 4140c578e6..ba7ae62da6 100644 --- a/automq-shell/src/main/java/com/automq/shell/stream/ClientStreamManager.java +++ b/automq-shell/src/main/java/com/automq/shell/stream/ClientStreamManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.stream; diff --git a/automq-shell/src/main/java/com/automq/shell/util/CLIUtils.java b/automq-shell/src/main/java/com/automq/shell/util/CLIUtils.java index b740066488..571be350b4 100644 --- a/automq-shell/src/main/java/com/automq/shell/util/CLIUtils.java +++ b/automq-shell/src/main/java/com/automq/shell/util/CLIUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.util; diff --git a/automq-shell/src/main/java/com/automq/shell/util/S3PropUtil.java b/automq-shell/src/main/java/com/automq/shell/util/S3PropUtil.java index 6b277273c2..f2b955cd05 100644 --- a/automq-shell/src/main/java/com/automq/shell/util/S3PropUtil.java +++ b/automq-shell/src/main/java/com/automq/shell/util/S3PropUtil.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.shell.util; diff --git a/automq-shell/src/main/java/com/automq/shell/util/Utils.java b/automq-shell/src/main/java/com/automq/shell/util/Utils.java new file mode 100644 index 0000000000..20975c1b1f --- /dev/null +++ b/automq-shell/src/main/java/com/automq/shell/util/Utils.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.shell.util; + +import com.automq.stream.s3.ByteBufAlloc; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +import io.netty.buffer.ByteBuf; + +public class Utils { + + public static ByteBuf compress(ByteBuf input) throws IOException { + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream); + + byte[] buffer = new byte[input.readableBytes()]; + input.readBytes(buffer); + gzipOutputStream.write(buffer); + gzipOutputStream.close(); + + ByteBuf compressed = ByteBufAlloc.byteBuffer(byteArrayOutputStream.size()); + compressed.writeBytes(byteArrayOutputStream.toByteArray()); + return compressed; + } + + public static ByteBuf decompress(ByteBuf input) throws IOException { + byte[] compressedData = new byte[input.readableBytes()]; + input.readBytes(compressedData); + ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(compressedData); + GZIPInputStream gzipInputStream = new GZIPInputStream(byteArrayInputStream); + + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = gzipInputStream.read(buffer)) != -1) { + byteArrayOutputStream.write(buffer, 0, bytesRead); + } + + gzipInputStream.close(); + byteArrayOutputStream.close(); + + byte[] uncompressedData = byteArrayOutputStream.toByteArray(); + ByteBuf output = ByteBufAlloc.byteBuffer(uncompressedData.length); + output.writeBytes(uncompressedData); + return output; + } +} diff --git a/automq-shell/src/main/resources/template/topo.yaml b/automq-shell/src/main/resources/template/topo.yaml index ca1f91933a..dd32d9d96d 100644 --- a/automq-shell/src/main/resources/template/topo.yaml +++ b/automq-shell/src/main/resources/template/topo.yaml @@ -9,10 +9,12 @@ global: config: | s3.data.buckets=0@s3://xxx_bucket?region=us-east-1 s3.ops.buckets=1@s3://xxx_bucket?region=us-east-1 + s3.wal.path=0@s3://xxx_bucket?region=us-east-1 + log.dirs=/root/kraft-logs envs: - - name: KAFKA_S3_ACCESS_KEY + - name: AWS_ACCESS_KEY_ID value: 'xxxxx' - - name: KAFKA_S3_SECRET_KEY + - name: AWS_SECRET_ACCESS_KEY value: 'xxxxx' controllers: # The controllers default are combined nodes which roles are controller and broker. diff --git a/automq-shell/src/test/java/com/automq/shell/util/UtilsTest.java b/automq-shell/src/test/java/com/automq/shell/util/UtilsTest.java new file mode 100644 index 0000000000..0398fc700a --- /dev/null +++ b/automq-shell/src/test/java/com/automq/shell/util/UtilsTest.java @@ -0,0 +1,50 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.shell.util; + +import com.automq.stream.s3.ByteBufAlloc; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import io.netty.buffer.ByteBuf; + +@Timeout(60) +@Tag("S3Unit") +public class UtilsTest { + + @Test + public void testCompression() { + String testStr = "This is a test string"; + ByteBuf input = ByteBufAlloc.byteBuffer(testStr.length()); + input.writeBytes(testStr.getBytes()); + try { + ByteBuf compressed = Utils.compress(input); + ByteBuf decompressed = Utils.decompress(compressed); + String decompressedStr = decompressed.toString(io.netty.util.CharsetUtil.UTF_8); + System.out.printf("Original: %s, Decompressed: %s\n", testStr, decompressedStr); + Assertions.assertEquals(testStr, decompressedStr); + } catch (Exception e) { + Assertions.fail("Exception occurred during compression/decompression: " + e.getMessage()); + } + } +} diff --git a/automq_release.py b/automq_release.py index 276222de8a..a9660d1dbf 100644 --- a/automq_release.py +++ b/automq_release.py @@ -70,7 +70,7 @@ def cmd(action, cmd_arg, *args, **kwargs): return print("*************************************************") - print("*** First command failure occurred here. ***") + print("*** First command failure occurred here. ***") print("*************************************************") fail("") diff --git a/bin/automq-perf-test.sh b/bin/automq-perf-test.sh index dd065326fe..aec6d203ec 100755 --- a/bin/automq-perf-test.sh +++ b/bin/automq-perf-test.sh @@ -23,4 +23,10 @@ fi if [ "x$KAFKA_HEAP_OPTS" = "x" ]; then export KAFKA_HEAP_OPTS="-Xmx1024M" fi +# Add additional help info for the new parameter (this won't be displayed directly but documents the change) +# --consumers-during-catchup: Percentage of consumers to activate during catch-up read (0-100, default: 100) +# This allows controlling what percentage of consumer groups are activated during catch-up +# reading to better simulate real-world scenarios where only a subset of consumers +# experience catch-up reads at the same time. + exec "$(dirname "$0")/kafka-run-class.sh" -name kafkaClient -loggc org.apache.kafka.tools.automq.PerfCommand "$@" diff --git a/bin/connect-distributed.sh b/bin/connect-distributed.sh index b8088ad923..62aafa9bc5 100755 --- a/bin/connect-distributed.sh +++ b/bin/connect-distributed.sh @@ -42,4 +42,5 @@ case $COMMAND in ;; esac +export KAFKA_CONNECT_MODE=true exec $(dirname $0)/kafka-run-class.sh $EXTRA_ARGS org.apache.kafka.connect.cli.ConnectDistributed "$@" diff --git a/bin/connect-standalone.sh b/bin/connect-standalone.sh index bef78d658f..24087a18da 100755 --- a/bin/connect-standalone.sh +++ b/bin/connect-standalone.sh @@ -42,4 +42,5 @@ case $COMMAND in ;; esac +export KAFKA_CONNECT_MODE=true exec $(dirname $0)/kafka-run-class.sh $EXTRA_ARGS org.apache.kafka.connect.cli.ConnectStandalone "$@" diff --git a/bin/kafka-run-class.sh b/bin/kafka-run-class.sh index fa70ded289..51d6065292 100755 --- a/bin/kafka-run-class.sh +++ b/bin/kafka-run-class.sh @@ -40,7 +40,23 @@ should_include_file() { fi file=$1 if [ -z "$(echo "$file" | grep -E "$regex")" ] ; then - return 0 + # If Connect mode is enabled, apply additional filtering + if [ "$KAFKA_CONNECT_MODE" = "true" ]; then + # Skip if file doesn't exist + [ ! -f "$file" ] && return 1 + + # Exclude heavy dependencies that Connect doesn't need + case "$file" in + *hadoop*) return 1 ;; + *hive*) return 1 ;; + *iceberg*) return 1 ;; + *avro*) return 1 ;; + *parquet*) return 1 ;; + *) return 0 ;; + esac + else + return 0 + fi else return 1 fi @@ -329,9 +345,9 @@ if [ "x$GC_LOG_ENABLED" = "xtrue" ]; then # We need to match to the end of the line to prevent sed from printing the characters that do not match # JAVA_MAJOR_VERSION=$("$JAVA" -version 2>&1 | sed -E -n 's/.* version "([0-9]*).*$/\1/p') if [[ "$JAVA_MAJOR_VERSION" -ge "9" ]] ; then - KAFKA_GC_LOG_OPTS="-Xlog:gc*:file=$LOG_DIR/$GC_LOG_FILE_NAME:time,tags:filecount=10,filesize=100M" + KAFKA_GC_LOG_OPTS="-Xlog:gc*:file=$LOG_DIR/$GC_LOG_FILE_NAME:time,tags:filecount=32,filesize=32M" else - KAFKA_GC_LOG_OPTS="-Xloggc:$LOG_DIR/$GC_LOG_FILE_NAME -verbose:gc -XX:+PrintGCDetails -XX:+PrintGCDateStamps -XX:+PrintGCTimeStamps -XX:+UseGCLogFileRotation -XX:NumberOfGCLogFiles=10 -XX:GCLogFileSize=100M" + KAFKA_GC_LOG_OPTS="-Xloggc:$LOG_DIR/$GC_LOG_FILE_NAME -verbose:gc -XX:+PrintGCDetails -XX:+PrintGCDateStamps -XX:+PrintGCTimeStamps -XX:+UseGCLogFileRotation -XX:NumberOfGCLogFiles=32 -XX:GCLogFileSize=32M" fi fi diff --git a/bin/kafka-server-start.sh b/bin/kafka-server-start.sh index 38702e9ee3..d235fa4e40 100755 --- a/bin/kafka-server-start.sh +++ b/bin/kafka-server-start.sh @@ -47,7 +47,7 @@ if [ "x$KAFKA_HEAP_OPTS" = "x" ]; then fi if [ "x$KAFKA_OPTS" = "x" ]; then - export KAFKA_OPTS="-Dio.netty.allocator.maxOrder=11" + export KAFKA_OPTS="-XX:+ExitOnOutOfMemoryError -XX:+HeapDumpOnOutOfMemoryError -Dio.netty.allocator.maxOrder=11" fi EXTRA_ARGS=${EXTRA_ARGS-'-name kafkaServer -loggc'} diff --git a/build.gradle b/build.gradle index ed7f763691..5470533eb4 100644 --- a/build.gradle +++ b/build.gradle @@ -53,7 +53,7 @@ plugins { ext { gradleVersion = versions.gradle - minJavaVersion = 11 + minJavaVersion = 17 buildVersionFileName = "kafka-version.properties" defaultMaxHeapSize = "2g" @@ -128,6 +128,9 @@ allprojects { repositories { mavenCentral() + maven { + url = uri("https://packages.confluent.io/maven/") + } } dependencyUpdates { @@ -147,6 +150,10 @@ allprojects { } configurations.all { + // Globally exclude commons-logging and logback to ensure a single logging implementation (reload4j) + exclude group: "commons-logging", module: "commons-logging" + exclude group: "ch.qos.logback", module: "logback-classic" + exclude group: "ch.qos.logback", module: "logback-core" // zinc is the Scala incremental compiler, it has a configuration for its own dependencies // that are unrelated to the project dependencies, we should not change them if (name != "zinc") { @@ -257,7 +264,10 @@ subprojects { options.compilerArgs << "-Xlint:-rawtypes" options.compilerArgs << "-Xlint:-serial" options.compilerArgs << "-Xlint:-try" - options.compilerArgs << "-Werror" + // AutoMQ inject start + // TODO: remove me, when upgrade to 4.x +// options.compilerArgs << "-Werror" + // AutoMQ inject start // --release is the recommended way to select the target release, but it's only supported in Java 9 so we also // set --source and --target via `sourceCompatibility` and `targetCompatibility` a couple of lines below @@ -728,7 +738,7 @@ subprojects { jacoco { toolVersion = versions.jacoco } - + jacocoTestReport { dependsOn tasks.test sourceSets sourceSets.main @@ -752,8 +762,8 @@ subprojects { skipProjects = [ ":jmh-benchmarks", ":trogdor" ] skipConfigurations = [ "zinc" ] } - // the task `removeUnusedImports` is implemented by google-java-format, - // and unfortunately the google-java-format version used by spotless 6.14.0 can't work with JDK 21. + // the task `removeUnusedImports` is implemented by google-java-format, + // and unfortunately the google-java-format version used by spotless 6.14.0 can't work with JDK 21. // Hence, we apply spotless tasks only if the env is either JDK11 or JDK17 if ((JavaVersion.current().isJava11() || (JavaVersion.current() == JavaVersion.VERSION_17))) { apply plugin: 'com.diffplug.spotless' @@ -828,6 +838,13 @@ tasks.create(name: "jarConnect", dependsOn: connectPkgs.collect { it + ":jar" }) tasks.create(name: "testConnect", dependsOn: connectPkgs.collect { it + ":test" }) {} +// OpenTelemetry related tasks +tasks.create(name: "jarOpenTelemetry", dependsOn: ":opentelemetry:jar") {} + +tasks.create(name: "testOpenTelemetry", dependsOn: ":opentelemetry:test") {} + +tasks.create(name: "buildOpenTelemetry", dependsOn: [":opentelemetry:jar", ":opentelemetry:test"]) {} + project(':server') { base { archivesName = "kafka-server" @@ -928,6 +945,8 @@ project(':core') { implementation project(':storage') implementation project(':server') implementation project(':automq-shell') + implementation project(':automq-metrics') + implementation project(':automq-log-uploader') implementation libs.argparse4j implementation libs.commonsValidator @@ -965,17 +984,77 @@ project(':core') { implementation libs.guava implementation libs.slf4jBridge implementation libs.slf4jReload4j + // The `jcl-over-slf4j` library is used to redirect JCL logging to SLF4J. + implementation libs.jclOverSlf4j - implementation libs.opentelemetryJava8 - implementation libs.opentelemetryOshi - implementation libs.opentelemetrySdk - implementation libs.opentelemetrySdkMetrics - implementation libs.opentelemetryExporterLogging - implementation libs.opentelemetryExporterProm - implementation libs.opentelemetryExporterOTLP - implementation libs.opentelemetryJmx implementation libs.awsSdkAuth + // table topic start + implementation ("org.apache.avro:avro:${versions.avro}") + implementation ("org.apache.avro:avro-protobuf:${versions.avro}") + implementation('com.google.protobuf:protobuf-java:3.25.5') + implementation ("org.apache.iceberg:iceberg-core:${versions.iceberg}") + implementation ("org.apache.iceberg:iceberg-api:${versions.iceberg}") + implementation ("org.apache.iceberg:iceberg-data:${versions.iceberg}") + implementation ("org.apache.iceberg:iceberg-parquet:${versions.iceberg}") + implementation ("org.apache.iceberg:iceberg-common:${versions.iceberg}") + implementation ("org.apache.iceberg:iceberg-aws:${versions.iceberg}") + implementation ("org.apache.iceberg:iceberg-nessie:${versions.iceberg}") + implementation ("software.amazon.awssdk:glue:${versions.awsSdk}") + implementation ("software.amazon.awssdk:s3tables:${versions.awsSdk}") + implementation 'software.amazon.s3tables:s3-tables-catalog-for-iceberg:0.1.0' + + implementation ('org.apache.hadoop:hadoop-common:3.4.1') { + exclude group: 'org.eclipse.jetty', module: '*' + exclude group: 'com.sun.jersey', module: '*' + } + // for hadoop common + implementation ("org.eclipse.jetty:jetty-webapp:${versions.jetty}") + + implementation (libs.kafkaAvroSerializer) { + exclude group: 'org.apache.kafka', module: 'kafka-clients' + } + + // > hive ext start + implementation 'org.apache.iceberg:iceberg-hive-metastore:1.6.1' + implementation('org.apache.hive:hive-metastore:3.1.3') { + // Remove useless dependencies (copy from iceberg-kafka-connect) + exclude group: "org.apache.avro", module: "avro" + exclude group: "org.slf4j", module: "slf4j-log4j12" + exclude group: "org.pentaho" // missing dependency + exclude group: "org.apache.hbase" + exclude group: "org.apache.logging.log4j" + exclude group: "co.cask.tephra" + exclude group: "com.google.code.findbugs", module: "jsr305" + exclude group: "org.eclipse.jetty.aggregate", module: "jetty-all" + exclude group: "org.eclipse.jetty.orbit", module: "javax.servlet" + exclude group: "org.apache.parquet", module: "parquet-hadoop-bundle" + exclude group: "com.tdunning", module: "json" + exclude group: "javax.transaction", module: "transaction-api" + exclude group: "com.zaxxer", module: "HikariCP" + exclude group: "org.apache.hadoop", module: "hadoop-yarn-server-common" + exclude group: "org.apache.hadoop", module: "hadoop-yarn-server-applicationhistoryservice" + exclude group: "org.apache.hadoop", module: "hadoop-yarn-server-resourcemanager" + exclude group: "org.apache.hadoop", module: "hadoop-yarn-server-web-proxy" + exclude group: "org.apache.hive", module: "hive-service-rpc" + exclude group: "com.github.joshelser", module: "dropwizard-metrics-hadoop-metrics2-reporter" + } + implementation ('org.apache.hadoop:hadoop-mapreduce-client-core:3.4.1') { + exclude group: 'com.sun.jersey', module: '*' + exclude group: 'com.sun.jersey.contribs', module: '*' + exclude group: 'com.github.pjfanning', module: 'jersey-json' + } + // > hive ext end + + // > Protobuf ext start + // Wire Runtime for schema handling + implementation ("com.squareup.wire:wire-schema:${versions.wire}") + implementation ("com.squareup.wire:wire-runtime:${versions.wire}") + implementation 'com.google.api.grpc:proto-google-common-protos:2.52.0' + // > Protobuf ext end + + // table topic end + implementation(libs.oshi) { exclude group: 'org.slf4j', module: '*' } @@ -990,6 +1069,7 @@ project(':core') { testImplementation project(':storage:storage-api').sourceSets.test.output testImplementation project(':server').sourceSets.test.output testImplementation libs.bcpkix + testImplementation libs.mockitoJunitJupiter // supports MockitoExtension testImplementation libs.mockitoCore testImplementation libs.guava testImplementation(libs.apacheda) { @@ -1160,7 +1240,6 @@ project(':core') { from(project.file("$rootDir/docker/docker-compose.yaml")) { into "docker/" } from(project.file("$rootDir/docker/telemetry")) { into "docker/telemetry/" } from(project.file("$rootDir/LICENSE")) { into "" } - from(project.file("$rootDir/LICENSE.S3Stream")) { into "" } from "$rootDir/NOTICE-binary" rename {String filename -> filename.replace("-binary", "")} from(configurations.runtimeClasspath) { into("libs/") } from(configurations.archives.artifacts.files) { into("libs/") } @@ -1171,6 +1250,10 @@ project(':core') { from(project(':trogdor').configurations.runtimeClasspath) { into("libs/") } from(project(':automq-shell').jar) { into("libs/") } from(project(':automq-shell').configurations.runtimeClasspath) { into("libs/") } + from(project(':automq-metrics').jar) { into("libs/") } + from(project(':automq-metrics').configurations.runtimeClasspath) { into("libs/") } + from(project(':automq-log-uploader').jar) { into("libs/") } + from(project(':automq-log-uploader').configurations.runtimeClasspath) { into("libs/") } from(project(':shell').jar) { into("libs/") } from(project(':shell').configurations.runtimeClasspath) { into("libs/") } from(project(':connect:api').jar) { into("libs/") } @@ -1202,6 +1285,38 @@ project(':core') { duplicatesStrategy 'exclude' } + // AutoMQ inject start + tasks.create(name: "releaseE2ETar", dependsOn: [configurations.archives.artifacts, 'copyDependantTestLibs'], type: Tar) { + def prefix = project.findProperty('prefix') ?: '' + archiveBaseName = "${prefix}kafka" + + into "${prefix}kafka-${archiveVersion.get()}" + compression = Compression.GZIP + from(project.file("$rootDir/bin")) { into "bin/" } + from(project.file("$rootDir/config")) { into "config/" } + from(project.file("$rootDir/licenses")) { into "licenses/" } + from(project.file("$rootDir/docker/docker-compose.yaml")) { into "docker/" } + from(project.file("$rootDir/docker/telemetry")) { into "docker/telemetry/" } + from(project.file("$rootDir/LICENSE")) { into "" } + from "$rootDir/NOTICE-binary" rename {String filename -> filename.replace("-binary", "")} + from(configurations.runtimeClasspath) { into("libs/") } + from(configurations.archives.artifacts.files) { into("libs/") } + from(project.siteDocsTar) { into("site-docs/") } + + // Include main and test jars from all subprojects + rootProject.subprojects.each { subproject -> + if (subproject.tasks.findByName('jar')) { + from(subproject.tasks.named('jar')) { into('libs/') } + } + if (subproject.tasks.findByName('testJar')) { + from(subproject.tasks.named('testJar')) { into('libs/') } + } + from(subproject.configurations.runtimeClasspath) { into('libs/') } + } + duplicatesStrategy 'exclude' + } + // AutoMQ inject end + jar { dependsOn('copyDependantLibs') } @@ -1220,7 +1335,7 @@ project(':core') { //By default gradle does not handle test dependencies between the sub-projects //This line is to include clients project test jar to dependant-testlibs from (project(':clients').testJar ) { "$buildDir/dependant-testlibs" } - // log4j-appender is not in core dependencies, + // log4j-appender is not in core dependencies, // so we add it to dependant-testlibs to avoid ClassNotFoundException in running kafka_log4j_appender.py from (project(':log4j-appender').jar ) { "$buildDir/dependant-testlibs" } duplicatesStrategy 'exclude' @@ -1253,6 +1368,7 @@ project(':core') { } } + project(':metadata') { base { archivesName = "kafka-metadata" @@ -1275,6 +1391,7 @@ project(':metadata') { implementation libs.guava implementation libs.awsSdkAuth implementation project(':s3stream') + implementation ("org.apache.avro:avro:${versions.avro}") implementation libs.jacksonDatabind implementation libs.jacksonJDK8Datatypes @@ -1480,7 +1597,7 @@ project(':transaction-coordinator') { implementation project(':clients') generator project(':generator') } - + sourceSets { main { java { @@ -1845,7 +1962,12 @@ project(':server-common') { implementation libs.jacksonDatabind implementation libs.pcollections implementation libs.opentelemetrySdk + + // AutoMQ inject start implementation project(':s3stream') + implementation libs.commonLang + // AutoMQ inject end + testImplementation project(':clients') testImplementation project(':clients').sourceSets.test.output @@ -2140,11 +2262,12 @@ project(':s3stream') { implementation 'commons-codec:commons-codec:1.17.0' implementation 'org.hdrhistogram:HdrHistogram:2.2.2' implementation 'software.amazon.awssdk.crt:aws-crt:0.30.8' + implementation 'com.ibm.async:asyncutil:0.1.0' - testImplementation 'org.slf4j:slf4j-simple:2.0.9' - testImplementation 'org.junit.jupiter:junit-jupiter:5.10.0' - testImplementation 'org.mockito:mockito-core:5.5.0' - testImplementation 'org.mockito:mockito-junit-jupiter:5.5.0' + testImplementation 'org.slf4j:slf4j-simple:1.7.36' + testImplementation libs.junitJupiter + testImplementation libs.mockitoCore + testImplementation libs.mockitoJunitJupiter // supports MockitoExtension testImplementation 'org.awaitility:awaitility:4.2.1' } @@ -2211,6 +2334,100 @@ project(':tools:tools-api') { } } +project(':automq-metrics') { + archivesBaseName = "automq-metrics" + + checkstyle { + configProperties = checkstyleConfigProperties("import-control-server.xml") + } + + dependencies { + // OpenTelemetry core dependencies + api libs.opentelemetryJava8 + api libs.opentelemetryOshi + api libs.opentelemetrySdk + api libs.opentelemetrySdkMetrics + api libs.opentelemetryExporterLogging + api libs.opentelemetryExporterProm + api libs.opentelemetryExporterOTLP + api libs.opentelemetryJmx + + // Logging dependencies + api libs.slf4jApi + api libs.slf4jBridge // 添加 SLF4J Bridge 依赖 + api libs.reload4j + + api libs.commonLang + + // Yammer metrics (for integration) + api 'com.yammer.metrics:metrics-core:2.2.0' + + implementation(project(':s3stream')) { + exclude(group: 'io.opentelemetry', module: '*') + exclude(group: 'io.opentelemetry.instrumentation', module: '*') + exclude(group: 'io.opentelemetry.proto', module: '*') + exclude(group: 'io.netty', module: 'netty-tcnative-boringssl-static') + exclude(group: 'com.github.jnr', module: '*') + exclude(group: 'org.aspectj', module: '*') + exclude(group: 'net.java.dev.jna', module: '*') + exclude(group: 'net.sourceforge.argparse4j', module: '*') + exclude(group: 'com.bucket4j', module: '*') + exclude(group: 'com.yammer.metrics', module: '*') + exclude(group: 'com.github.spotbugs', module: '*') + exclude(group: 'org.apache.kafka.shaded', module: '*') + } + implementation libs.nettyBuffer + implementation libs.jacksonDatabind + implementation libs.guava + implementation project(':clients') + + // Test dependencies + testImplementation libs.junitJupiter + testImplementation libs.mockitoCore + testImplementation libs.slf4jReload4j + + testRuntimeOnly libs.junitPlatformLanucher + + implementation('io.opentelemetry:opentelemetry-sdk:1.40.0') + implementation("io.opentelemetry.semconv:opentelemetry-semconv:1.25.0-alpha") + implementation("io.opentelemetry.instrumentation:opentelemetry-runtime-telemetry-java8:2.6.0-alpha") + implementation('com.google.protobuf:protobuf-java:3.25.5') + implementation('org.xerial.snappy:snappy-java:1.1.10.5') + } + + clean.doFirst { + delete "$buildDir/kafka/" + } + + javadoc { + enabled = false + } +} + +project(':automq-log-uploader') { + archivesBaseName = "automq-log-uploader" + + checkstyle { + configProperties = checkstyleConfigProperties("import-control-server.xml") + } + + dependencies { + api project(':s3stream') + + implementation project(':clients') + implementation libs.reload4j + implementation libs.slf4jApi + implementation libs.slf4jBridge + implementation libs.nettyBuffer + implementation libs.guava + implementation libs.commonLang + } + + javadoc { + enabled = false + } +} + project(':tools') { base { archivesName = "kafka-tools" @@ -2226,7 +2443,16 @@ project(':tools') { implementation (project(':log4j-appender')){ exclude group: 'org.slf4j', module: '*' } + // AutoMQ inject start implementation project(':automq-shell') + implementation (libs.kafkaAvroSerializer) { + exclude group: 'org.apache.kafka', module: 'kafka-clients' + } + implementation libs.bucket4j + implementation (libs.oshi){ + exclude group: 'org.slf4j', module: '*' + } + // AutoMQ inject end implementation project(':storage') implementation project(':connect:runtime') @@ -2242,6 +2468,7 @@ project(':tools') { implementation libs.awsSdkAuth implementation libs.hdrHistogram implementation libs.spotbugsAnnotations + implementation libs.guava // for SASL/OAUTHBEARER JWT validation implementation (libs.jose4j){ @@ -3308,6 +3535,8 @@ project(':connect:runtime') { api project(':clients') api project(':connect:json') api project(':connect:transforms') + api project(':automq-metrics') + api project(':automq-log-uploader') implementation libs.slf4jApi implementation libs.reload4j @@ -3316,6 +3545,7 @@ project(':connect:runtime') { implementation libs.jacksonJaxrsJsonProvider implementation libs.jerseyContainerServlet implementation libs.jerseyHk2 + implementation libs.jaxrsApi implementation libs.jaxbApi // Jersey dependency that was available in the JDK before Java 9 implementation libs.activation // Jersey dependency that was available in the JDK before Java 9 implementation libs.jettyServer diff --git a/chart/bitnami/README.md b/chart/bitnami/README.md new file mode 100644 index 0000000000..e7ea20867b --- /dev/null +++ b/chart/bitnami/README.md @@ -0,0 +1,62 @@ +# AutoMQ + +[AutoMQ](https://www.automq.com/) is a cloud-native alternative to Kafka by decoupling durability to cloud storage services like S3. 10x Cost-Effective. No Cross-AZ Traffic Cost. Autoscale in seconds. Single-digit ms latency. +This Helm chart simplifies the deployment of AutoMQ into your Kubernetes cluster using the Software model. + +## Prerequisites +### Install Helm chart +Install Helm chart and version v3.8.0+ +[Helm chart quickstart](https://helm.sh/zh/docs/intro/quickstart/) +```shell +helm version +``` +### Using the Bitnami Helm repository +AutoMQ is fully compatible with Bitnami's Helm Charts, so you can customize your AutoMQ Kubernetes cluster based on the relevant values.yaml of Bitnami. +[Bitnami Helm Charts](https://github.com/bitnami/charts) + +## Quickstart +### Setup a Kubernetes Cluster +The quickest way to set up a Kubernetes cluster to install Bitnami Charts is by following the "Bitnami Get Started" guides for the different services: + +[Get Started with Bitnami Charts using the Amazon Elastic Container Service for Kubernetes (EKS)](https://docs.bitnami.com/kubernetes/get-started-eks/) + + +### Installing the AutoMQ with Bitnami Chart + +As an alternative to supplying the configuration parameters as arguments, you can create a supplemental YAML file containing your specific config parameters. Any parameters not specified in this file will default to those set in [values.yaml](values.yaml). + +1. Create an empty `automq-values.yaml` file +2. Edit the file with your specific parameters: + + You can refer to the [demo-values.yaml](/chart/bitnami/demo-values.yaml) based on the bitnami [values.yaml](https://github.com/bitnami/charts/blob/main/bitnami/kafka/values.yaml) + we provided for deploying AutoMQ on AWS across 3 Availability Zones using m7g.xlarge instances (4 vCPUs, 16GB Mem, 156MiB/s network bandwidth). + + + You need to replace the bucket configurations in the placeholders $｛｝, such as ops-bucket, data-bucket, region, endpoint, access-key/secret-key. + +3. Install or upgrade the AutoMQ Helm chart using your custom yaml file: + + we recommend using the `--version` [31.x.x (31.1.0 ~ 31.5.0)](https://artifacthub.io/packages/helm/bitnami/kafka) bitnami helm chart while installing AutoMQ. + +```shell +helm install automq-release oci://registry-1.docker.io/bitnamicharts/kafka -f demo-values.yaml --version 31.5.0 --namespace automq --create-namespace +``` + +### Upgrading + +To upgrade the deployment: + +```shell +helm repo update +helm upgrade automq-release oci://registry-1.docker.io/bitnamicharts/kafka -f demo-values.yaml --version 31.5.0 --namespace automq --create-namespace +``` + +### Uninstalling the Chart + +To uninstall/delete the deployment: + +```shell +helm uninstall automq-release --namespace automq +``` + +This command removes all the Kubernetes components associated with the chart and deletes the release. diff --git a/chart/bitnami/demo-values.yaml b/chart/bitnami/demo-values.yaml new file mode 100644 index 0000000000..3004ba895e --- /dev/null +++ b/chart/bitnami/demo-values.yaml @@ -0,0 +1,141 @@ +global: + security: + allowInsecureImages: true +image: + registry: automqinc + repository: automq + tag: 1.5.0-bitnami + pullPolicy: Always +extraEnvVars: + - name: AWS_ACCESS_KEY_ID + value: "${access-key}" + - name: AWS_SECRET_ACCESS_KEY + value: "${secret-key}" +controller: + replicaCount: 3 + resources: + requests: + cpu: "3000m" + memory: "12Gi" + limits: + cpu: "4000m" + memory: "16Gi" + heapOpts: -Xmx6g -Xms6g -XX:MaxDirectMemorySize=6g -XX:MetaspaceSize=96m + extraConfig: | + elasticstream.enable=true + autobalancer.client.auth.sasl.mechanism=PLAIN + autobalancer.client.auth.sasl.jaas.config=org.apache.kafka.common.security.plain.PlainLoginModule required username="inter_broker_user" password="interbroker-password-placeholder" user_inter_broker_user="interbroker-password-placeholder"; + autobalancer.client.auth.security.protocol=SASL_PLAINTEXT + autobalancer.client.listener.name=INTERNAL + s3.wal.cache.size=2147483648 + s3.block.cache.size=1073741824 + s3.stream.allocator.policy=POOLED_DIRECT + s3.network.baseline.bandwidth=245366784 + # Replace the following with your bucket config + s3.ops.buckets=1@s3://${ops-bucket}?region=${region}&endpoint=${endpoint} + s3.data.buckets=0@s3://${data-bucket}?region=${region}&endpoint=${endpoint} + s3.wal.path=0@s3://${data-bucket}?region=${region}&endpoint=${endpoint} + automq.zonerouter.channels=0@s3://${data-bucket}?region=${region}&endpoint=${endpoint} + affinity: + podAntiAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + - labelSelector: + matchExpressions: + - key: app.kubernetes.io/instance + operator: In + # your helm release name + values: + - automq-release + - key: app.kubernetes.io/component + operator: In + values: + - controller-eligible + - broker + topologyKey: kubernetes.io/hostname +# --- nodeAffinity recommended --- +# nodeAffinity: +# requiredDuringSchedulingIgnoredDuringExecution: +# nodeSelectorTerms: +# - matchExpressions: +# - key: "${your-node-label-key}" +# operator: In +# values: +# - "${your-node-label-value}" + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: DoNotSchedule + labelSelector: + matchLabels: + app.kubernetes.io/component: controller-eligible + tolerations: + - key: "dedicated" + operator: "Equal" + value: "automq" + effect: "NoSchedule" + persistence: + size: 20Gi + +broker: + replicaCount: 3 + resources: + requests: + cpu: "3000m" + memory: "12Gi" + limits: + cpu: "4000m" + memory: "16Gi" + heapOpts: -Xmx6g -Xms6g -XX:MaxDirectMemorySize=6g -XX:MetaspaceSize=96m + extraConfig: | + elasticstream.enable=true + autobalancer.client.auth.sasl.mechanism=PLAIN + autobalancer.client.auth.sasl.jaas.config=org.apache.kafka.common.security.plain.PlainLoginModule required username="inter_broker_user" password="interbroker-password-placeholder" user_inter_broker_user="interbroker-password-placeholder"; + autobalancer.client.auth.security.protocol=SASL_PLAINTEXT + autobalancer.client.listener.name=INTERNAL + s3.wal.cache.size=2147483648 + s3.block.cache.size=1073741824 + s3.stream.allocator.policy=POOLED_DIRECT + s3.network.baseline.bandwidth=245366784 + # Replace the following with your bucket config + s3.ops.buckets=1@s3://${ops-bucket}?region=${region}&endpoint=${endpoint} + s3.data.buckets=0@s3://${data-bucket}?region=${region}&endpoint=${endpoint} + s3.wal.path=0@s3://${data-bucket}?region=${region}&endpoint=${endpoint} + automq.zonerouter.channels=0@s3://${data-bucket}?region=${region}&endpoint=${endpoint} + affinity: + podAntiAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + - labelSelector: + matchExpressions: + - key: app.kubernetes.io/instance + operator: In + # your helm release name + values: + - automq-release + - key: app.kubernetes.io/component + operator: In + values: + - controller-eligible + - broker + topologyKey: kubernetes.io/hostname +# --- nodeAffinity recommended --- +# nodeAffinity: +# requiredDuringSchedulingIgnoredDuringExecution: +# nodeSelectorTerms: +# - matchExpressions: +# - key: "${your-node-label-key}" +# operator: In +# values: +# - "${your-node-label-value}" + topologySpreadConstraints: + - maxSkew: 1 + topologyKey: topology.kubernetes.io/zone + whenUnsatisfiable: DoNotSchedule + labelSelector: + matchLabels: + app.kubernetes.io/component: broker + tolerations: + - key: "dedicated" + operator: "Equal" + value: "automq" + effect: "NoSchedule" +brokerRackAssignment: aws-az diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index e1930fa0a2..2a2ae4ff82 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -372,11 +372,12 @@ + files="(StreamControlManager|S3StreamsMetadataImage|CompactionManagerTest|S3StreamMetricsManager|CompactionManager|BlockCache|DefaultS3BlockCache|StreamReader|S3Utils|AnomalyDetector|Recreate|ForceClose|QuorumController|AbstractObjectStorage).java"/> + diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java index a4f73d5d2e..a5e4e910b3 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java @@ -1729,6 +1729,16 @@ default ListClientMetricsResourcesResult listClientMetricsResources() { * @return {@link GetNodesResult} */ GetNodesResult getNodes(Collection nodeIdList, GetNodesOptions options); + + /** + * Update consumer group + * + * @param groupId group id + * @param groupSpec {@link UpdateGroupSpec} + * @param options {@link UpdateGroupOptions} + * @return {@link UpdateGroupResult} + */ + UpdateGroupResult updateGroup(String groupId, UpdateGroupSpec groupSpec, UpdateGroupOptions options); // AutoMQ inject end /** diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ForwardingAdmin.java b/clients/src/main/java/org/apache/kafka/clients/admin/ForwardingAdmin.java index 6f86dbcf54..8822553259 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/ForwardingAdmin.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ForwardingAdmin.java @@ -314,5 +314,11 @@ public RemoveRaftVoterResult removeRaftVoter(int voterId, Uuid voterDirectoryId, public GetNodesResult getNodes(Collection nodeIdList, GetNodesOptions options) { return delegate.getNodes(nodeIdList, options); } + + @Override + public UpdateGroupResult updateGroup(String groupId, UpdateGroupSpec groupSpec, UpdateGroupOptions options) { + return delegate.updateGroup(groupId, groupSpec, options); + } + // AutoMQ inject end } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesOptions.java index 43633c50ec..833c097c76 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesOptions.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.clients.admin; diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesResult.java index 0cce564f0e..47824386be 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesResult.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/GetNodesResult.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.clients.admin; diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java index 88e7a74d2f..12202776ac 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java @@ -56,6 +56,7 @@ import org.apache.kafka.clients.admin.internals.ListOffsetsHandler; import org.apache.kafka.clients.admin.internals.ListTransactionsHandler; import org.apache.kafka.clients.admin.internals.RemoveMembersFromConsumerGroupHandler; +import org.apache.kafka.clients.admin.internals.UpdateGroupHandler; import org.apache.kafka.clients.consumer.OffsetAndMetadata; import org.apache.kafka.clients.consumer.internals.ConsumerProtocol; import org.apache.kafka.common.Cluster; @@ -4857,6 +4858,14 @@ void handleFailure(Throwable throwable) { return new GetNodesResult(future); } + @Override + public UpdateGroupResult updateGroup(String groupId, UpdateGroupSpec groupSpec, UpdateGroupOptions options) { + SimpleAdminApiFuture future = UpdateGroupHandler.newFuture(groupId); + UpdateGroupHandler handler = new UpdateGroupHandler(groupId, groupSpec, logContext); + invokeDriver(handler, future, options.timeoutMs); + return new UpdateGroupResult(future.get(CoordinatorKey.byGroupId(groupId))); + } + private void invokeDriver( AdminApiHandler handler, AdminApiFuture future, diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/NodeMetadata.java b/clients/src/main/java/org/apache/kafka/clients/admin/NodeMetadata.java index bcc4895768..dac5cfb6e9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/NodeMetadata.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/NodeMetadata.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.clients.admin; diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupOptions.java new file mode 100644 index 0000000000..c5f5c2ec37 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupOptions.java @@ -0,0 +1,23 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +public class UpdateGroupOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupResult.java new file mode 100644 index 0000000000..83469f3275 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupResult.java @@ -0,0 +1,37 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.KafkaFuture; + +public class UpdateGroupResult extends AbstractOptions { + private final KafkaFuture future; + + UpdateGroupResult(final KafkaFuture future) { + this.future = future; + } + + /** + * Return a future which succeeds if all the feature updates succeed. + */ + public KafkaFuture all() { + return future; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupSpec.java b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupSpec.java new file mode 100644 index 0000000000..4b548313ae --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateGroupSpec.java @@ -0,0 +1,68 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import java.util.Objects; + +public class UpdateGroupSpec { + private String linkId; + private boolean promoted; + + public UpdateGroupSpec linkId(String linkId) { + this.linkId = linkId; + return this; + } + + public UpdateGroupSpec promoted(boolean promoted) { + this.promoted = promoted; + return this; + } + + public String linkId() { + return linkId; + } + + public boolean promoted() { + return promoted; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (o == null || getClass() != o.getClass()) + return false; + UpdateGroupSpec spec = (UpdateGroupSpec) o; + return promoted == spec.promoted && Objects.equals(linkId, spec.linkId); + } + + @Override + public int hashCode() { + return Objects.hash(linkId, promoted); + } + + @Override + public String toString() { + return "UpdateGroupsSpec{" + + "linkId='" + linkId + '\'' + + ", promoted=" + promoted + + '}'; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/UpdateGroupHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/UpdateGroupHandler.java new file mode 100644 index 0000000000..384b3ba422 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/UpdateGroupHandler.java @@ -0,0 +1,144 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.admin.internals; + +import org.apache.kafka.clients.admin.UpdateGroupSpec; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.message.AutomqUpdateGroupRequestData; +import org.apache.kafka.common.message.AutomqUpdateGroupResponseData; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.AbstractResponse; +import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType; +import org.apache.kafka.common.requests.s3.AutomqUpdateGroupRequest; +import org.apache.kafka.common.requests.s3.AutomqUpdateGroupResponse; +import org.apache.kafka.common.utils.LogContext; + +import org.slf4j.Logger; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static java.util.Collections.singleton; + +public class UpdateGroupHandler extends AdminApiHandler.Batched { + private final CoordinatorKey groupId; + private final UpdateGroupSpec groupSpec; + private final Logger logger; + private final AdminApiLookupStrategy lookupStrategy; + + public UpdateGroupHandler( + String groupId, + UpdateGroupSpec groupSpec, + LogContext logContext + ) { + this.groupId = CoordinatorKey.byGroupId(groupId); + this.groupSpec = groupSpec; + this.logger = logContext.logger(UpdateGroupHandler.class); + this.lookupStrategy = new CoordinatorStrategy(CoordinatorType.GROUP, logContext); + } + + @Override + public String apiName() { + return "updateGroup"; + } + + @Override + public AdminApiLookupStrategy lookupStrategy() { + return lookupStrategy; + } + + public static AdminApiFuture.SimpleAdminApiFuture newFuture( + String groupId + ) { + return AdminApiFuture.forKeys(Collections.singleton(CoordinatorKey.byGroupId(groupId))); + } + + private void validateKeys(Set groupIds) { + if (!groupIds.equals(singleton(groupId))) { + throw new IllegalArgumentException("Received unexpected group ids " + groupIds + + " (expected only " + singleton(groupId) + ")"); + } + } + + @Override + public AutomqUpdateGroupRequest.Builder buildBatchedRequest( + int coordinatorId, + Set groupIds + ) { + validateKeys(groupIds); + return new AutomqUpdateGroupRequest.Builder( + new AutomqUpdateGroupRequestData() + .setLinkId(groupSpec.linkId()) + .setGroupId(this.groupId.idValue) + .setPromoted(groupSpec.promoted()) + ); + } + + @Override + public ApiResult handleResponse( + Node coordinator, + Set groupIds, + AbstractResponse abstractResponse + ) { + validateKeys(groupIds); + + final Map completed = new HashMap<>(); + final Map failed = new HashMap<>(); + final List groupsToUnmap = new ArrayList<>(); + AutomqUpdateGroupResponse response = (AutomqUpdateGroupResponse) abstractResponse; + AutomqUpdateGroupResponseData data = response.data(); + Errors error = Errors.forCode(data.errorCode()); + if (error != Errors.NONE) { + handleError( + CoordinatorKey.byGroupId(data.groupId()), + error, + data.errorMessage(), + failed, + groupsToUnmap + ); + } else { + completed.put(groupId, null); + } + return new ApiResult<>(completed, failed, groupsToUnmap); + } + + private void handleError( + CoordinatorKey groupId, + Errors error, + String errorMsg, + Map failed, + List groupsToUnmap + ) { + switch (error) { + case COORDINATOR_NOT_AVAILABLE: + case NOT_COORDINATOR: + // If the coordinator is unavailable or there was a coordinator change, then we unmap + // the key so that we retry the `FindCoordinator` request + logger.debug("`{}` request for group id {} returned error {}. " + + "Will attempt to find the coordinator again and retry.", apiName(), groupId.idValue, error); + groupsToUnmap.add(groupId); + break; + default: + logger.error("`{}` request for group id {} failed due to unexpected error {}.", apiName(), groupId.idValue, error); + failed.put(groupId, error.exception(errorMsg)); + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java b/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java index 3689227d1f..33be0bf052 100755 --- a/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java +++ b/clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java @@ -81,17 +81,17 @@ public class TopicConfig { public static final String REMOTE_LOG_STORAGE_ENABLE_CONFIG = "remote.storage.enable"; public static final String REMOTE_LOG_STORAGE_ENABLE_DOC = "To enable tiered storage for a topic, set this configuration as true. " + - "You can not disable this config once it is enabled. It will be provided in future versions."; + "You can not disable this config once it is enabled. It will be provided in future versions."; public static final String LOCAL_LOG_RETENTION_MS_CONFIG = "local.retention.ms"; public static final String LOCAL_LOG_RETENTION_MS_DOC = "The number of milliseconds to keep the local log segment before it gets deleted. " + - "Default value is -2, it represents `retention.ms` value is to be used. The effective value should always be less than or equal " + - "to `retention.ms` value."; + "Default value is -2, it represents `retention.ms` value is to be used. The effective value should always be less than or equal " + + "to `retention.ms` value."; public static final String LOCAL_LOG_RETENTION_BYTES_CONFIG = "local.retention.bytes"; public static final String LOCAL_LOG_RETENTION_BYTES_DOC = "The maximum size of local log segments that can grow for a partition before it " + - "deletes the old segments. Default value is -2, it represents `retention.bytes` value to be used. The effective value should always be " + - "less than or equal to `retention.bytes` value."; + "deletes the old segments. Default value is -2, it represents `retention.bytes` value to be used. The effective value should always be " + + "less than or equal to `retention.bytes` value."; public static final String REMOTE_LOG_DISABLE_POLICY_RETAIN = "retain"; public static final String REMOTE_LOG_DISABLE_POLICY_DELETE = "delete"; @@ -103,16 +103,16 @@ public class TopicConfig { "selected then all data in remote will be kept post-disablement and will only be deleted when it breaches expiration " + "thresholds. If %s is selected then the data will be made inaccessible immediately by advancing the log start offset and will be " + "deleted asynchronously.", REMOTE_LOG_DISABLE_POLICY_RETAIN, REMOTE_LOG_DISABLE_POLICY_DELETE, - REMOTE_LOG_DISABLE_POLICY_RETAIN, REMOTE_LOG_DISABLE_POLICY_DELETE); + REMOTE_LOG_DISABLE_POLICY_RETAIN, REMOTE_LOG_DISABLE_POLICY_DELETE); public static final String MAX_MESSAGE_BYTES_CONFIG = "max.message.bytes"; public static final String MAX_MESSAGE_BYTES_DOC = "The largest record batch size allowed by Kafka (after compression if compression is enabled). " + - "If this is increased and there are consumers older than 0.10.2, the consumers' fetch " + - "size must also be increased so that they can fetch record batches this large. " + - "In the latest message format version, records are always grouped into batches for efficiency. " + - "In previous message format versions, uncompressed records are not grouped into batches and this " + - "limit only applies to a single record in that case."; + "If this is increased and there are consumers older than 0.10.2, the consumers' fetch " + + "size must also be increased so that they can fetch record batches this large. " + + "In the latest message format version, records are always grouped into batches for efficiency. " + + "In previous message format versions, uncompressed records are not grouped into batches and this " + + "limit only applies to a single record in that case."; public static final String INDEX_INTERVAL_BYTES_CONFIG = "index.interval.bytes"; public static final String INDEX_INTERVAL_BYTES_DOC = "This setting controls how frequently " + @@ -256,4 +256,90 @@ public class TopicConfig { "broker will not perform down-conversion for consumers expecting an older message format. The broker responds " + "with UNSUPPORTED_VERSION error for consume requests from such older clients. This configuration" + "does not apply to any message format conversion that might be required for replication to followers."; + + // AutoMQ inject start + public static final String TABLE_TOPIC_ENABLE_CONFIG = "automq.table.topic.enable"; + public static final String TABLE_TOPIC_ENABLE_DOC = "The configuration controls whether enable table topic"; + public static final String TABLE_TOPIC_COMMIT_INTERVAL_CONFIG = "automq.table.topic.commit.interval.ms"; + public static final String TABLE_TOPIC_COMMIT_INTERVAL_DOC = "The table topic commit interval(ms)"; + public static final String TABLE_TOPIC_NAMESPACE_CONFIG = "automq.table.topic.namespace"; + public static final String TABLE_TOPIC_NAMESPACE_DOC = "The table topic table namespace"; + + public static final String TABLE_TOPIC_SCHEMA_TYPE_CONFIG = "automq.table.topic.schema.type"; + public static final String TABLE_TOPIC_SCHEMA_TYPE_DOC = "[DEPRECATED] The table topic schema type configuration. " + + "This configuration is deprecated and will be removed in a future release. " + + "Please use the new separate converter and transform configurations instead. " + + "Supported values: 'schemaless' (maps to convert.value.type=raw, transform.value.type=none), " + + "'schema' (maps to convert.value.type=by_schema_id, transform.value.type=flatten)."; + + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG = "automq.table.topic.convert.value.type"; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_DOC = "How to parse Kafka record values. " + + "Supported: 'raw', 'string', 'by_schema_id', 'by_latest_schema'. " + + "Schema Registry URL required for 'by_schema_id' and 'by_latest_schema'."; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG = "automq.table.topic.convert.key.type"; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_DOC = "How to parse Kafka record keys. " + + "Supported: 'raw', 'string', 'by_schema_id', 'by_latest_schema'. " + + "Schema Registry URL required for 'by_schema_id' and 'by_latest_schema'."; + + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_SUBJECT_CONFIG = + "automq.table.topic.convert.value.by_latest_schema.subject"; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_SUBJECT_DOC = + "Subject name to resolve the latest value schema from Schema Registry when using convert.value.type=by_latest_schema. " + + "If not set, defaults to '-value'."; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_CONFIG = + "automq.table.topic.convert.value.by_latest_schema.message.full.name"; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_DOC = + "Fully-qualified message name for the latest value schema (if using Protobuf) when convert.value.type=by_latest_schema." + + "If not set, uses the first message."; + + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_SUBJECT_CONFIG = + "automq.table.topic.convert.key.by_latest_schema.subject"; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_SUBJECT_DOC = + "Subject name to resolve the latest key schema from Schema Registry when using convert.key.type=by_latest_schema. " + + "If not set, defaults to '-key'."; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_CONFIG = + "automq.table.topic.convert.key.by_latest_schema.message.full.name"; + public static final String AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_DOC = + "Fully-qualified message name for the latest key schema (if using Protobuf) when convert.key.type=by_latest_schema. " + + "If not set, uses the first message."; + + public static final String AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG = "automq.table.topic.transform.value.type"; + public static final String AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_DOC = "Transformation to apply to the record value after conversion. " + + "Supported: 'none', 'flatten' (extract fields from structured records), " + + "'flatten_debezium' (process Debezium CDC events). " + + "Note: 'flatten_debezium' requires schema-based conversion."; + + public static final String TABLE_TOPIC_ID_COLUMNS_CONFIG = "automq.table.topic.id.columns"; + public static final String TABLE_TOPIC_ID_COLUMNS_DOC = "The primary key, comma-separated list of columns that identify a row in tables." + + "ex. [region, name]"; + public static final String TABLE_TOPIC_PARTITION_BY_CONFIG = "automq.table.topic.partition.by"; + public static final String TABLE_TOPIC_PARTITION_BY_DOC = "The partition fields of the table. ex. [bucket(name), month(timestamp)]"; + public static final String TABLE_TOPIC_UPSERT_ENABLE_CONFIG = "automq.table.topic.upsert.enable"; + public static final String TABLE_TOPIC_UPSERT_ENABLE_DOC = "The configuration controls whether enable table topic upsert"; + public static final String TABLE_TOPIC_CDC_FIELD_CONFIG = "automq.table.topic.cdc.field"; + public static final String TABLE_TOPIC_CDC_FIELD_DOC = "The name of the field containing the CDC operation, I, U, or D"; + + public static final String AUTOMQ_TABLE_TOPIC_ERRORS_TOLERANCE_CONFIG = "automq.table.topic.errors.tolerance"; + public static final String AUTOMQ_TABLE_TOPIC_ERRORS_TOLERANCE_DOC = "Configures the error handling strategy for table topic record processing. Valid values are none, invalid_data, and all."; + + public static final String AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_CONFIG = "automq.table.topic.expire.snapshot.enabled"; + public static final String AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_DOC = "Enable/disable automatic snapshot expiration."; + public static final boolean AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_DEFAULT = true; + + public static final String AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_CONFIG = "automq.table.topic.expire.snapshot.older.than.hours"; + public static final String AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_DOC = "Set retention duration in hours."; + public static final int AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_DEFAULT = 1; + + public static final String AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_CONFIG = "automq.table.topic.expire.snapshot.retain.last"; + public static final String AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_DOC = "Minimum snapshots to retain."; + public static final int AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_DEFAULT = 1; + + public static final String KAFKA_LINKS_ID_CONFIG = "automq.kafka.links.id"; + public static final String KAFKA_LINKS_ID_DOC = "The unique id of a kafka link"; + public static final String KAFKA_LINKS_TOPIC_START_TIME_CONFIG = "automq.kafka.links.topic.start.time"; + public static final String KAFKA_LINKS_TOPIC_START_TIME_DOC = "The offset to start replicate from. Valid values: -1 (latest), -2 (earliest), positive value (timestamp)"; + public static final String KAFKA_LINKS_TOPIC_STATE_CONFIG = "automq.kafka.links.topic.state"; + public static final String KAFKA_LINKS_TOPIC_STATE_DOC = "The state of the topic that's in linking"; + // AutoMQ inject end + } diff --git a/clients/src/main/java/org/apache/kafka/common/errors/s3/NodeLockedException.java b/clients/src/main/java/org/apache/kafka/common/errors/s3/NodeLockedException.java new file mode 100644 index 0000000000..e5b6f3c330 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/common/errors/s3/NodeLockedException.java @@ -0,0 +1,30 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.common.errors.s3; + +import org.apache.kafka.common.errors.ApiException; + +public class NodeLockedException extends ApiException { + + public NodeLockedException(String message) { + super(message); + } + +} diff --git a/clients/src/main/java/org/apache/kafka/common/errors/s3/ObjectNotCommittedException.java b/clients/src/main/java/org/apache/kafka/common/errors/s3/ObjectNotCommittedException.java new file mode 100644 index 0000000000..416ee06b15 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/common/errors/s3/ObjectNotCommittedException.java @@ -0,0 +1,28 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.common.errors.s3; + +import org.apache.kafka.common.errors.ApiException; + +public class ObjectNotCommittedException extends ApiException { + public ObjectNotCommittedException(String message) { + super(message); + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/errors/s3/UnregisterNodeWithOpenStreamsException.java b/clients/src/main/java/org/apache/kafka/common/errors/s3/UnregisterNodeWithOpenStreamsException.java index 280f8fed24..9681d7188e 100644 --- a/clients/src/main/java/org/apache/kafka/common/errors/s3/UnregisterNodeWithOpenStreamsException.java +++ b/clients/src/main/java/org/apache/kafka/common/errors/s3/UnregisterNodeWithOpenStreamsException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.errors.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/internals/Topic.java b/clients/src/main/java/org/apache/kafka/common/internals/Topic.java index d94c91da50..3709d8005b 100644 --- a/clients/src/main/java/org/apache/kafka/common/internals/Topic.java +++ b/clients/src/main/java/org/apache/kafka/common/internals/Topic.java @@ -30,7 +30,13 @@ public class Topic { public static final String TRANSACTION_STATE_TOPIC_NAME = "__transaction_state"; public static final String SHARE_GROUP_STATE_TOPIC_NAME = "__share_group_state"; public static final String CLUSTER_METADATA_TOPIC_NAME = "__cluster_metadata"; + + // AutoMQ inject start public static final String AUTO_BALANCER_METRICS_TOPIC_NAME = "__auto_balancer_metrics"; + public static final String TABLE_TOPIC_CONTROL_TOPIC_NAME = "__automq_table_control"; + public static final String TABLE_TOPIC_DATA_TOPIC_NAME = "__automq_table_data"; + // AutoMQ inject end + public static final TopicPartition CLUSTER_METADATA_TOPIC_PARTITION = new TopicPartition( CLUSTER_METADATA_TOPIC_NAME, 0 diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java index 1d31855db5..2d21150b21 100644 --- a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java +++ b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java @@ -129,4 +129,13 @@ public void config(MetricConfig config) { this.config = config; } } + + // AutoMQ inject start + /** + * A public method to expose the {@link #measurableValue} method. + */ + public double measurableValueV2(long timeMs) { + return measurableValue(timeMs); + } + // AutoMQ inject end } diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java b/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java index 92c5617ae1..aa36ee1b12 100644 --- a/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java +++ b/clients/src/main/java/org/apache/kafka/common/protocol/ApiKeys.java @@ -148,9 +148,11 @@ public enum ApiKeys { AUTOMQ_REGISTER_NODE(ApiMessageType.AUTOMQ_REGISTER_NODE, false, false), AUTOMQ_GET_NODES(ApiMessageType.AUTOMQ_GET_NODES, false, true), AUTOMQ_ZONE_ROUTER(ApiMessageType.AUTOMQ_ZONE_ROUTER, false, false), + AUTOMQ_GET_PARTITION_SNAPSHOT(ApiMessageType.AUTOMQ_GET_PARTITION_SNAPSHOT, false, false), GET_NEXT_NODE_ID(ApiMessageType.GET_NEXT_NODE_ID, false, true), - DESCRIBE_STREAMS(ApiMessageType.DESCRIBE_STREAMS, false, true); + DESCRIBE_STREAMS(ApiMessageType.DESCRIBE_STREAMS, false, true), + AUTOMQ_UPDATE_GROUP(ApiMessageType.AUTOMQ_UPDATE_GROUP); // AutoMQ for Kafka inject end private static final Map> APIS_BY_LISTENER = diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java index 35581c45c1..b886a24a1f 100644 --- a/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java +++ b/clients/src/main/java/org/apache/kafka/common/protocol/Errors.java @@ -151,6 +151,8 @@ import org.apache.kafka.common.errors.s3.NodeEpochExpiredException; import org.apache.kafka.common.errors.s3.NodeEpochNotExistException; import org.apache.kafka.common.errors.s3.NodeFencedException; +import org.apache.kafka.common.errors.s3.NodeLockedException; +import org.apache.kafka.common.errors.s3.ObjectNotCommittedException; import org.apache.kafka.common.errors.s3.ObjectNotExistException; import org.apache.kafka.common.errors.s3.OffsetNotMatchedException; import org.apache.kafka.common.errors.s3.RedundantOperationException; @@ -435,6 +437,8 @@ public enum Errors { KEY_EXIST(512, "The key already exists.", KeyExistException::new), KEY_NOT_EXIST(513, "The key does not exist.", ObjectNotExistException::new), NODE_FENCED(514, "The node is fenced.", NodeFencedException::new), + NODE_LOCKED(515, "The node is locked", NodeLockedException::new), + OBJECT_NOT_COMMITED(516, "The object is not commited.", ObjectNotCommittedException::new), STREAM_INNER_ERROR(599, "The stream inner error.", StreamInnerErrorException::new), // AutoMQ inject end diff --git a/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java index 9ab8715236..4a20c23ef7 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java +++ b/clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java @@ -499,6 +499,11 @@ public void setPartitionLeaderEpoch(int epoch) { throw new UnsupportedOperationException("Magic versions prior to 2 do not support partition leader epoch"); } + @Override + public void setProducerId(long producerId) { + throw new UnsupportedOperationException("Magic versions prior to 2 do not support producer id"); + } + private void setTimestampAndUpdateCrc(TimestampType timestampType, long timestamp) { byte attributes = LegacyRecord.computeAttributes(magic(), compressionType(), timestampType); buffer.put(LOG_OVERHEAD + LegacyRecord.ATTRIBUTES_OFFSET, attributes); diff --git a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java index 7d8ee75e19..ff1437f117 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java +++ b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java @@ -190,6 +190,16 @@ public long producerId() { return buffer.getLong(PRODUCER_ID_OFFSET); } + @Override + public void setProducerId(long producerId) { + if (producerId() == producerId) { + return; + } + buffer.putLong(PRODUCER_ID_OFFSET, producerId); + long crc = computeChecksum(); + ByteUtils.writeUnsignedInt(buffer, CRC_OFFSET, crc); + } + @Override public short producerEpoch() { return buffer.getShort(PRODUCER_EPOCH_OFFSET); diff --git a/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java b/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java index b5f42e5b91..f98f078912 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java +++ b/clients/src/main/java/org/apache/kafka/common/record/MutableRecordBatch.java @@ -65,4 +65,13 @@ public interface MutableRecordBatch extends RecordBatch { * @return The closeable iterator */ CloseableIterator skipKeyValueIterator(BufferSupplier bufferSupplier); + + // AutoMQ injection start + + /** + * Set the producer id for this batch of records. + * @param producerId The producer id to use + */ + void setProducerId(long producerId); + // AutoMQ injection end } diff --git a/clients/src/main/java/org/apache/kafka/common/record/PooledRecords.java b/clients/src/main/java/org/apache/kafka/common/record/PooledRecords.java index 57a5330dbb..9c75071f35 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/PooledRecords.java +++ b/clients/src/main/java/org/apache/kafka/common/record/PooledRecords.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.record; diff --git a/clients/src/main/java/org/apache/kafka/common/record/TimestampType.java b/clients/src/main/java/org/apache/kafka/common/record/TimestampType.java index becde9dc85..04baaa4530 100644 --- a/clients/src/main/java/org/apache/kafka/common/record/TimestampType.java +++ b/clients/src/main/java/org/apache/kafka/common/record/TimestampType.java @@ -39,6 +39,22 @@ public static TimestampType forName(String name) { throw new NoSuchElementException("Invalid timestamp type " + name); } + public static TimestampType forId(int id) { + switch (id) { + case -1: { + return NO_TIMESTAMP_TYPE; + } + case 0: { + return CREATE_TIME; + } + case 1: { + return LOG_APPEND_TIME; + } + default: + throw new IllegalArgumentException("Invalid timestamp type " + id); + } + } + @Override public String toString() { return name; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java index 9e50b558ef..1f5fe7e88a 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractRequest.java @@ -24,7 +24,9 @@ import org.apache.kafka.common.protocol.ObjectSerializationCache; import org.apache.kafka.common.protocol.SendBuilder; import org.apache.kafka.common.requests.s3.AutomqGetNodesRequest; +import org.apache.kafka.common.requests.s3.AutomqGetPartitionSnapshotRequest; import org.apache.kafka.common.requests.s3.AutomqRegisterNodeRequest; +import org.apache.kafka.common.requests.s3.AutomqUpdateGroupRequest; import org.apache.kafka.common.requests.s3.AutomqZoneRouterRequest; import org.apache.kafka.common.requests.s3.CloseStreamsRequest; import org.apache.kafka.common.requests.s3.CommitStreamObjectRequest; @@ -375,10 +377,14 @@ private static AbstractRequest doParseRequest(ApiKeys apiKey, short apiVersion, return AutomqGetNodesRequest.parse(buffer, apiVersion); case AUTOMQ_ZONE_ROUTER: return AutomqZoneRouterRequest.parse(buffer, apiVersion); + case AUTOMQ_GET_PARTITION_SNAPSHOT: + return AutomqGetPartitionSnapshotRequest.parse(buffer, apiVersion); case GET_NEXT_NODE_ID: return GetNextNodeIdRequest.parse(buffer, apiVersion); case DESCRIBE_STREAMS: return DescribeStreamsRequest.parse(buffer, apiVersion); + case AUTOMQ_UPDATE_GROUP: + return AutomqUpdateGroupRequest.parse(buffer, apiVersion); // AutoMQ for Kafka inject end case SHARE_GROUP_HEARTBEAT: diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java index f9952c9eed..545c077be5 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractResponse.java @@ -22,7 +22,9 @@ import org.apache.kafka.common.protocol.MessageUtil; import org.apache.kafka.common.protocol.SendBuilder; import org.apache.kafka.common.requests.s3.AutomqGetNodesResponse; +import org.apache.kafka.common.requests.s3.AutomqGetPartitionSnapshotResponse; import org.apache.kafka.common.requests.s3.AutomqRegisterNodeResponse; +import org.apache.kafka.common.requests.s3.AutomqUpdateGroupResponse; import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; import org.apache.kafka.common.requests.s3.CloseStreamsResponse; import org.apache.kafka.common.requests.s3.CommitStreamObjectResponse; @@ -312,10 +314,14 @@ public static AbstractResponse parseResponse(ApiKeys apiKey, ByteBuffer response return AutomqGetNodesResponse.parse(responseBuffer, version); case AUTOMQ_ZONE_ROUTER: return AutomqZoneRouterResponse.parse(responseBuffer, version); + case AUTOMQ_GET_PARTITION_SNAPSHOT: + return AutomqGetPartitionSnapshotResponse.parse(responseBuffer, version); case GET_NEXT_NODE_ID: return GetNextNodeIdResponse.parse(responseBuffer, version); case DESCRIBE_STREAMS: return DescribeStreamsResponse.parse(responseBuffer, version); + case AUTOMQ_UPDATE_GROUP: + return AutomqUpdateGroupResponse.parse(responseBuffer, version); // AutoMQ for Kafka inject end case SHARE_GROUP_HEARTBEAT: diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesRequest.java index 8276944127..afece35873 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesRequest.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesResponse.java index 9e8ed65f96..cc31bbd5cd 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetNodesResponse.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetPartitionSnapshotRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetPartitionSnapshotRequest.java new file mode 100644 index 0000000000..f23ae1ceed --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetPartitionSnapshotRequest.java @@ -0,0 +1,76 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.common.requests.s3; + +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotRequestData; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.requests.AbstractRequest; +import org.apache.kafka.common.requests.AbstractResponse; +import org.apache.kafka.common.requests.ApiError; + +import java.nio.ByteBuffer; + +public class AutomqGetPartitionSnapshotRequest extends AbstractRequest { + private final AutomqGetPartitionSnapshotRequestData data; + + public AutomqGetPartitionSnapshotRequest(AutomqGetPartitionSnapshotRequestData data, short version) { + super(ApiKeys.AUTOMQ_GET_PARTITION_SNAPSHOT, version); + this.data = data; + } + + @Override + public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) { + ApiError apiError = ApiError.fromThrowable(e); + AutomqGetPartitionSnapshotResponseData response = new AutomqGetPartitionSnapshotResponseData() + .setErrorCode(apiError.error().code()) + .setThrottleTimeMs(throttleTimeMs); + return new AutomqGetPartitionSnapshotResponse(response); + } + + @Override + public AutomqGetPartitionSnapshotRequestData data() { + return data; + } + + public static AutomqGetPartitionSnapshotRequest parse(ByteBuffer buffer, short version) { + return new AutomqGetPartitionSnapshotRequest(new AutomqGetPartitionSnapshotRequestData(new ByteBufferAccessor(buffer), version), version); + } + + public static class Builder extends AbstractRequest.Builder { + private final AutomqGetPartitionSnapshotRequestData data; + + public Builder(AutomqGetPartitionSnapshotRequestData data) { + super(ApiKeys.AUTOMQ_GET_PARTITION_SNAPSHOT); + this.data = data; + } + + @Override + public AutomqGetPartitionSnapshotRequest build(short version) { + return new AutomqGetPartitionSnapshotRequest(data, version); + } + + @Override + public String toString() { + return data.toString(); + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetPartitionSnapshotResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetPartitionSnapshotResponse.java new file mode 100644 index 0000000000..ac2a323473 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqGetPartitionSnapshotResponse.java @@ -0,0 +1,62 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.common.requests.s3; + +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.AbstractResponse; + +import java.nio.ByteBuffer; +import java.util.Map; + +public class AutomqGetPartitionSnapshotResponse extends AbstractResponse { + private final AutomqGetPartitionSnapshotResponseData data; + + public AutomqGetPartitionSnapshotResponse(AutomqGetPartitionSnapshotResponseData data) { + super(ApiKeys.AUTOMQ_GET_PARTITION_SNAPSHOT); + this.data = data; + } + + @Override + public Map errorCounts() { + return errorCounts(Errors.forCode(data.errorCode())); + } + + @Override + public int throttleTimeMs() { + return data.throttleTimeMs(); + } + + @Override + public void maybeSetThrottleTimeMs(int throttleTimeMs) { + data.setThrottleTimeMs(throttleTimeMs); + } + + @Override + public AutomqGetPartitionSnapshotResponseData data() { + return data; + } + + public static AutomqGetPartitionSnapshotResponse parse(ByteBuffer buffer, short version) { + return new AutomqGetPartitionSnapshotResponse(new AutomqGetPartitionSnapshotResponseData(new ByteBufferAccessor(buffer), version)); + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeRequest.java index 3d0b7383a0..8a62cc5bcd 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeRequest.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeResponse.java index df593a28a3..04c10a8038 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqRegisterNodeResponse.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqUpdateGroupRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqUpdateGroupRequest.java new file mode 100644 index 0000000000..5da79f1399 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqUpdateGroupRequest.java @@ -0,0 +1,77 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.common.requests.s3; + +import org.apache.kafka.common.message.AutomqUpdateGroupRequestData; +import org.apache.kafka.common.message.AutomqUpdateGroupResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.requests.AbstractRequest; +import org.apache.kafka.common.requests.AbstractResponse; +import org.apache.kafka.common.requests.ApiError; + +import java.nio.ByteBuffer; + +public class AutomqUpdateGroupRequest extends AbstractRequest { + private final AutomqUpdateGroupRequestData data; + + public AutomqUpdateGroupRequest(AutomqUpdateGroupRequestData data, short version) { + super(ApiKeys.AUTOMQ_UPDATE_GROUP, version); + this.data = data; + } + + @Override + public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) { + ApiError apiError = ApiError.fromThrowable(e); + AutomqUpdateGroupResponseData response = new AutomqUpdateGroupResponseData() + .setErrorCode(apiError.error().code()) + .setThrottleTimeMs(throttleTimeMs); + return new AutomqUpdateGroupResponse(response); + } + + @Override + public AutomqUpdateGroupRequestData data() { + return data; + } + + public static AutomqUpdateGroupRequest parse(ByteBuffer buffer, short version) { + return new AutomqUpdateGroupRequest(new AutomqUpdateGroupRequestData(new ByteBufferAccessor(buffer), version), version); + } + + public static class Builder extends AbstractRequest.Builder { + + private final AutomqUpdateGroupRequestData data; + + public Builder(AutomqUpdateGroupRequestData data) { + super(ApiKeys.AUTOMQ_UPDATE_GROUP); + this.data = data; + } + + @Override + public AutomqUpdateGroupRequest build(short version) { + return new AutomqUpdateGroupRequest(data, version); + } + + @Override + public String toString() { + return data.toString(); + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqUpdateGroupResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqUpdateGroupResponse.java new file mode 100644 index 0000000000..0b790bac59 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqUpdateGroupResponse.java @@ -0,0 +1,62 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.common.requests.s3; + +import org.apache.kafka.common.message.AutomqUpdateGroupResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.AbstractResponse; + +import java.nio.ByteBuffer; +import java.util.Map; + +public class AutomqUpdateGroupResponse extends AbstractResponse { + private final AutomqUpdateGroupResponseData data; + + public AutomqUpdateGroupResponse(AutomqUpdateGroupResponseData data) { + super(ApiKeys.AUTOMQ_UPDATE_GROUP); + this.data = data; + } + + @Override + public Map errorCounts() { + return errorCounts(Errors.forCode(data.errorCode())); + } + + @Override + public int throttleTimeMs() { + return data.throttleTimeMs(); + } + + @Override + public void maybeSetThrottleTimeMs(int throttleTimeMs) { + data.setThrottleTimeMs(throttleTimeMs); + } + + @Override + public AutomqUpdateGroupResponseData data() { + return data; + } + + public static AutomqUpdateGroupResponse parse(ByteBuffer buffer, short version) { + return new AutomqUpdateGroupResponse(new AutomqUpdateGroupResponseData(new ByteBufferAccessor(buffer), version)); + } +} diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterRequest.java index cf12961950..8b9ea02c2f 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterRequest.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterResponse.java index 5e3dae11d1..a9cd6be56e 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/AutomqZoneRouterResponse.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsRequest.java index c12a2f95a4..868755aae4 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsRequest.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsResponse.java index 7f24b2b023..f49339566a 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/s3/DescribeStreamsResponse.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.requests.s3; diff --git a/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java index 5b0828a082..d5f8b27ced 100644 --- a/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java +++ b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java @@ -23,5 +23,9 @@ public enum ClientQuotaType { PRODUCE, FETCH, REQUEST, + // AutoMQ for Kafka inject start + SLOW_FETCH, + REQUEST_RATE, + // AutoMQ for Kafka inject end CONTROLLER_MUTATION } diff --git a/clients/src/main/resources/common/message/AutomqGetPartitionSnapshotRequest.json b/clients/src/main/resources/common/message/AutomqGetPartitionSnapshotRequest.json new file mode 100644 index 0000000000..3d77731a7d --- /dev/null +++ b/clients/src/main/resources/common/message/AutomqGetPartitionSnapshotRequest.json @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +{ + "apiKey": 516, + "type": "request", + "listeners": [ + "broker" + ], + "name": "AutomqGetPartitionSnapshotRequest", + "validVersions": "0-2", + "flexibleVersions": "0+", + "fields": [ + { + "name": "SessionId", + "type": "int32", + "versions": "0+", + "about": "The get session id" + }, + { + "name": "SessionEpoch", + "type": "int32", + "versions": "0+", + "about": "The get session epoch, which is used for ordering requests in a session" + }, + { + "name": "RequestCommit", + "type": "bool", + "versions": "1+", + "about": "Request commit the ConfirmWAL data to the main storage." + }, + { + "name": "Version", + "type": "int16", + "versions": "1+", + "about": "The route request version" + } + ] +} \ No newline at end of file diff --git a/clients/src/main/resources/common/message/AutomqGetPartitionSnapshotResponse.json b/clients/src/main/resources/common/message/AutomqGetPartitionSnapshotResponse.json new file mode 100644 index 0000000000..cf48c61b21 --- /dev/null +++ b/clients/src/main/resources/common/message/AutomqGetPartitionSnapshotResponse.json @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +{ + "apiKey": 516, + "type": "response", + "name": "AutomqGetPartitionSnapshotResponse", + "validVersions": "0-2", + "flexibleVersions": "0+", + "fields": [ + { "name": "ErrorCode", "type": "int16", "versions": "0+", "about": "The top level response error code" }, + { "name": "ThrottleTimeMs", "type": "int32", "versions": "0+", "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." }, + { "name": "SessionId", "type": "int32", "versions": "0+", "about": "The get session id" }, + { "name": "SessionEpoch", "type": "int32", "versions": "0+", "about": "The next get session epoch" }, + { "name": "Topics", "type": "[]Topic", "versions": "0+", "about": "The topic list", "fields": [ + { "name": "TopicId", "type": "uuid", "versions": "0+", "about": "The topic id", "mapKey": true }, + { "name": "Partitions", "type": "[]PartitionSnapshot", "versions": "0+", "about": "The partition snapshot list", "fields": [ + { "name": "PartitionIndex", "type": "int32", "versions": "0+", "about": "The partition index"}, + { "name": "LeaderEpoch", "type": "int32", "versions": "0+", "about": "The partition leader epoch"}, + { "name": "Operation", "type": "int16", "versions": "0+", "about": "The snapshot operation, 0 -> ADD, 1 -> PATCH, 2 -> REMOVE"}, + { "name": "LogMetadata", "type": "LogMetadata", "versions": "0+", "nullableVersions": "0+", "about": "The log metadata" }, + { "name": "FirstUnstableOffset", "type": "LogOffsetMetadata", "versions": "0+", "nullableVersions": "0+", "about": "The partition first unstable offset" }, + { "name": "LogEndOffset", "type": "LogOffsetMetadata", "versions": "0+", "nullableVersions": "0+", "about": "The partition log end offset" }, + { "name": "StreamMetadata", "type": "[]StreamMetadata", "versions": "0+", "nullableVersions": "0+", "fields": [ + { "name": "StreamId", "type": "int64", "versions": "0+", "about": "The streamId" }, + { "name": "EndOffset", "type": "int64", "versions": "0+", "about": "The stream end offset" } + ]}, + { "name": "LastTimestampOffset", "type": "TimestampOffsetData", "versions": "1+", "nullableVersions": "1+", "about": "The last segment's last time index" } + ]} + ]}, + { + "name": "ConfirmWalEndOffset", + "type": "bytes", + "versions": "1+", + "about": "The confirm WAL end offset." + }, + { + "name": "ConfirmWalConfig", + "type": "string", + "versions": "1+", + "about": "The confirm WAL config." + }, + { + "name": "ConfirmWalDeltaData", + "type": "bytes", + "versions": "2+", + "nullableVersions": "2+", + "about": "The confirm WAL delta data between two end offsets. It's an optional field. If not present, the client should read the delta from WAL" + } + ], + "commonStructs": [ + { "name": "LogMetadata", "versions": "0+", "fields": [ + { "name": "segments", "versions": "0+", "type": "[]SegmentMetadata", "about": "The segment list" }, + { "name": "streamMap", "versions": "0+", "type": "[]StreamMapping" } + ]}, + { "name": "StreamMapping", "versions": "0+", "fields": [ + { "name": "name", "type": "string", "versions": "0+", "about": "The streamName", "mapKey": true }, + { "name": "streamId", "type": "int64", "versions": "0+", "about": "The stream id" } + ]}, + { "name": "LogOffsetMetadata", "versions": "0+", "fields": [ + { "name": "messageOffset", "type": "int64", "versions": "0+", "about": "The message logic offset" }, + { "name": "relativePositionInSegment", "type": "int32", "versions": "0+", "about": "The message relative physical offset" } + ]}, + { "name": "SegmentMetadata", "versions": "0+", "fields": [ + { "name": "baseOffset", "versions": "0+", "type": "int64", "about": "The segment base offset" }, + { "name": "createTimestamp", "versions": "0+", "type": "int64", "about": "The segment create timestamp" }, + { "name": "lastModifiedTimestamp", "versions": "0+", "type": "int64", "about": "The segment last modified timestamp" }, + { "name": "streamSuffix", "versions": "0+", "type": "string", "about": "The segment's stream suffix" }, + { "name": "logSize", "versions": "0+", "type": "int32", "about": "The segment size" }, + { "name": "log", "versions": "0+", "type": "SliceRange", "about": "The segment log stream slice range" }, + { "name": "time", "versions": "0+", "type": "SliceRange", "about": "The segment time stream slice range" }, + { "name": "transaction", "versions": "0+", "type": "SliceRange", "about": "The segment transaction stream slice range" }, + { "name": "firstBatchTimestamp", "versions": "0+", "type": "int64", "about": "The segment first batch timestamp" }, + { "name": "timeIndexLastEntry", "versions": "0+", "type": "TimestampOffsetData", "about": "The segment last timestamp index entry" } + ]}, + { "name": "SliceRange", "versions": "0+", "fields": [ + { "name": "start", "versions": "0+", "type": "int64", "about": "The range start offset" }, + { "name": "end", "versions": "0+", "type": "int64", "about": "The range end offset" } + ]}, + { "name": "TimestampOffsetData", "versions": "0+", "fields": [ + { "name": "timestamp", "versions": "0+", "type": "int64", "about": "The range start offset" }, + { "name": "offset", "versions": "0+", "type": "int64", "about": "The range end offset" } + ]} + ] +} \ No newline at end of file diff --git a/clients/src/main/resources/common/message/AutomqUpdateGroupRequest.json b/clients/src/main/resources/common/message/AutomqUpdateGroupRequest.json new file mode 100644 index 0000000000..8443159aa4 --- /dev/null +++ b/clients/src/main/resources/common/message/AutomqUpdateGroupRequest.json @@ -0,0 +1,45 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +{ + "apiKey": 602, + "type": "request", + "listeners": [ + "broker" + ], + "name": "AutomqUpdateGroupRequest", + "validVersions": "0", + "flexibleVersions": "0+", + "fields": [ + { + "name": "LinkId", + "type": "string", + "versions": "0+", + "about": "The unique ID of a kafka link" + }, + { + "name": "GroupId", + "type": "string", + "versions": "0+", + "about": "The group identifier." + }, + { + "name": "Promoted", + "type": "bool", + "versions": "0+", + "about": "Whether the group is promoted" + } + ] +} \ No newline at end of file diff --git a/clients/src/main/resources/common/message/AutomqUpdateGroupResponse.json b/clients/src/main/resources/common/message/AutomqUpdateGroupResponse.json new file mode 100644 index 0000000000..fa927cd419 --- /dev/null +++ b/clients/src/main/resources/common/message/AutomqUpdateGroupResponse.json @@ -0,0 +1,50 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +{ + "apiKey": 602, + "type": "response", + "name": "AutomqUpdateGroupResponse", + "validVersions": "0", + "flexibleVersions": "0+", + "fields": [ + { + "name": "GroupId", + "type": "string", + "versions": "0+", + "about": "The group identifier." + }, + { + "name": "ErrorCode", + "type": "int16", + "versions": "0+", + "about": "The error code, or 0 if there was no error." + }, + { + "name": "ErrorMessage", + "type": "string", + "versions": "0+", + "nullableVersions": "0+", + "default": "null", + "about": "The top-level error message, or null if there was no error." + }, + { + "name": "ThrottleTimeMs", + "type": "int32", + "versions": "0+", + "about": "Duration in milliseconds for which the request was throttled due to a quota violation, or zero if the request did not violate any quota." + } + ] +} \ No newline at end of file diff --git a/clients/src/main/resources/common/message/AutomqZoneRouterRequest.json b/clients/src/main/resources/common/message/AutomqZoneRouterRequest.json index 00586559bf..0797d97509 100644 --- a/clients/src/main/resources/common/message/AutomqZoneRouterRequest.json +++ b/clients/src/main/resources/common/message/AutomqZoneRouterRequest.json @@ -20,7 +20,7 @@ "broker" ], "name": "AutomqZoneRouterRequest", - "validVersions": "0", + "validVersions": "0-1", "flexibleVersions": "0+", "fields": [ { @@ -28,6 +28,18 @@ "type": "bytes", "versions": "0+", "about": "The router metadata" + }, + { + "name": "RouteEpoch", + "type": "int64", + "versions": "1+", + "about": "The route requests epoch" + }, + { + "name": "Version", + "type": "int16", + "versions": "1+", + "about": "The route request version" } ] } \ No newline at end of file diff --git a/clients/src/main/resources/common/message/AutomqZoneRouterResponse.json b/clients/src/main/resources/common/message/AutomqZoneRouterResponse.json index f3dc178eca..5dbf956817 100644 --- a/clients/src/main/resources/common/message/AutomqZoneRouterResponse.json +++ b/clients/src/main/resources/common/message/AutomqZoneRouterResponse.json @@ -17,7 +17,7 @@ "apiKey": 515, "type": "response", "name": "AutomqZoneRouterResponse", - "validVersions": "0", + "validVersions": "0-1", "flexibleVersions": "0+", "fields": [ { diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java b/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java index 7f2c0f8e5a..aea08dcd34 100644 --- a/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java +++ b/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java @@ -1444,6 +1444,12 @@ public synchronized Node broker(int index) { public GetNodesResult getNodes(Collection nodeIdList, GetNodesOptions options) { throw new UnsupportedOperationException(); } + + @Override + public UpdateGroupResult updateGroup(String groupId, UpdateGroupSpec groupSpec, UpdateGroupOptions options) { + throw new UnsupportedOperationException(); + } + // AutoMQ inject end } diff --git a/clients/src/test/java/org/apache/kafka/common/utils/GrowableMultiBufferSupplierTest.java b/clients/src/test/java/org/apache/kafka/common/utils/GrowableMultiBufferSupplierTest.java index d9e9d23cbc..a3071ad415 100644 --- a/clients/src/test/java/org/apache/kafka/common/utils/GrowableMultiBufferSupplierTest.java +++ b/clients/src/test/java/org/apache/kafka/common/utils/GrowableMultiBufferSupplierTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common.utils; diff --git a/config/connect-log4j.properties b/config/connect-log4j.properties index 61b2ac331d..35c4ae95f0 100644 --- a/config/connect-log4j.properties +++ b/config/connect-log4j.properties @@ -24,7 +24,8 @@ log4j.appender.stdout.layout=org.apache.log4j.PatternLayout # location of the log files (e.g. ${kafka.logs.dir}/connect.log). The `MaxFileSize` option specifies the maximum size of the log file, # and the `MaxBackupIndex` option specifies the number of backup files to keep. # -log4j.appender.connectAppender=org.apache.log4j.RollingFileAppender +log4j.appender.connectAppender=com.automq.log.S3RollingFileAppender +log4j.appender.connectAppender.configProviderClass=org.apache.kafka.connect.automq.log.ConnectS3LogConfigProvider log4j.appender.connectAppender.MaxFileSize=10MB log4j.appender.connectAppender.MaxBackupIndex=11 log4j.appender.connectAppender.File=${kafka.logs.dir}/connect.log diff --git a/config/kraft/broker.properties b/config/kraft/broker.properties index c0199b10a1..c941b8257f 100644 --- a/config/kraft/broker.properties +++ b/config/kraft/broker.properties @@ -143,8 +143,8 @@ s3.data.buckets=0@s3://ko3?region=us-east-1 # The ops buckets s3.ops.buckets=0@s3://ko3?region=us-east-1 -# The file path of delta WAL in block device -s3.wal.path=0@file:///tmp/kraft-broker-logs/s3wal?capacity=2147483648 +# The wal storage config +s3.wal.path=0@s3://ko3?region=us-east-1 # The maximum size of WAL cache can use, default 2GB # s3.wal.cache.size=2147483648 diff --git a/config/kraft/server.properties b/config/kraft/server.properties index 18606a3548..2741cc7909 100644 --- a/config/kraft/server.properties +++ b/config/kraft/server.properties @@ -146,8 +146,8 @@ s3.data.buckets=0@s3://ko3?region=us-east-1 # The ops buckets s3.ops.buckets=0@s3://ko3?region=us-east-1 -# The file path of delta WAL in block device -s3.wal.path=0@file:///tmp/kraft-broker-logs/s3wal?capacity=2147483648 +# The wal storage config +s3.wal.path=0@s3://ko3?region=us-east-1 # The maximum size of WAL cache can use, default 2GB # s3.wal.cache.size=2147483648 diff --git a/config/log4j.properties b/config/log4j.properties index db7b679bf1..fd67a3a98d 100644 --- a/config/log4j.properties +++ b/config/log4j.properties @@ -21,70 +21,73 @@ log4j.appender.stdout=org.apache.log4j.ConsoleAppender log4j.appender.stdout.layout=org.apache.log4j.PatternLayout log4j.appender.stdout.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.kafkaAppender=com.automq.shell.log.S3RollingFileAppender +log4j.logger.com.automq.log.S3RollingFileAppender=INFO, stdout +log4j.additivity.com.automq.log.S3RollingFileAppender=false + +log4j.appender.kafkaAppender=com.automq.log.S3RollingFileAppender log4j.appender.kafkaAppender.MaxFileSize=100MB log4j.appender.kafkaAppender.MaxBackupIndex=14 log4j.appender.kafkaAppender.File=${kafka.logs.dir}/server.log log4j.appender.kafkaAppender.layout=org.apache.log4j.PatternLayout log4j.appender.kafkaAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.stateChangeAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.stateChangeAppender=com.automq.log.S3RollingFileAppender log4j.appender.stateChangeAppender.MaxFileSize=10MB log4j.appender.stateChangeAppender.MaxBackupIndex=11 log4j.appender.stateChangeAppender.File=${kafka.logs.dir}/state-change.log log4j.appender.stateChangeAppender.layout=org.apache.log4j.PatternLayout log4j.appender.stateChangeAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.requestAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.requestAppender=com.automq.log.S3RollingFileAppender log4j.appender.requestAppender.MaxFileSize=10MB log4j.appender.requestAppender.MaxBackupIndex=11 log4j.appender.requestAppender.File=${kafka.logs.dir}/kafka-request.log log4j.appender.requestAppender.layout=org.apache.log4j.PatternLayout log4j.appender.requestAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.cleanerAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.cleanerAppender=com.automq.log.S3RollingFileAppender log4j.appender.cleanerAppender.MaxFileSize=10MB log4j.appender.cleanerAppender.MaxBackupIndex=11 log4j.appender.cleanerAppender.File=${kafka.logs.dir}/log-cleaner.log log4j.appender.cleanerAppender.layout=org.apache.log4j.PatternLayout log4j.appender.cleanerAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.controllerAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.controllerAppender=com.automq.log.S3RollingFileAppender log4j.appender.controllerAppender.MaxFileSize=100MB log4j.appender.controllerAppender.MaxBackupIndex=14 log4j.appender.controllerAppender.File=${kafka.logs.dir}/controller.log log4j.appender.controllerAppender.layout=org.apache.log4j.PatternLayout log4j.appender.controllerAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.authorizerAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.authorizerAppender=com.automq.log.S3RollingFileAppender log4j.appender.authorizerAppender.MaxFileSize=10MB log4j.appender.authorizerAppender.MaxBackupIndex=11 log4j.appender.authorizerAppender.File=${kafka.logs.dir}/kafka-authorizer.log log4j.appender.authorizerAppender.layout=org.apache.log4j.PatternLayout log4j.appender.authorizerAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.s3ObjectAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.s3ObjectAppender=com.automq.log.S3RollingFileAppender log4j.appender.s3ObjectAppender.MaxFileSize=100MB log4j.appender.s3ObjectAppender.MaxBackupIndex=14 log4j.appender.s3ObjectAppender.File=${kafka.logs.dir}/s3-object.log log4j.appender.s3ObjectAppender.layout=org.apache.log4j.PatternLayout log4j.appender.s3ObjectAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.s3StreamMetricsAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.s3StreamMetricsAppender=com.automq.log.S3RollingFileAppender log4j.appender.s3StreamMetricsAppender.MaxFileSize=10MB log4j.appender.s3StreamMetricsAppender.MaxBackupIndex=11 log4j.appender.s3StreamMetricsAppender.File=${kafka.logs.dir}/s3stream-metrics.log log4j.appender.s3StreamMetricsAppender.layout=org.apache.log4j.PatternLayout log4j.appender.s3StreamMetricsAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.s3StreamThreadPoolAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.s3StreamThreadPoolAppender=com.automq.log.S3RollingFileAppender log4j.appender.s3StreamThreadPoolAppender.MaxFileSize=10MB log4j.appender.s3StreamThreadPoolAppender.MaxBackupIndex=11 log4j.appender.s3StreamThreadPoolAppender.File=${kafka.logs.dir}/s3stream-threads.log log4j.appender.s3StreamThreadPoolAppender.layout=org.apache.log4j.PatternLayout log4j.appender.s3StreamThreadPoolAppender.layout.ConversionPattern=[%d] %p %m (%c)%n -log4j.appender.autoBalancerAppender=com.automq.shell.log.S3RollingFileAppender +log4j.appender.autoBalancerAppender=com.automq.log.S3RollingFileAppender log4j.appender.autoBalancerAppender.MaxFileSize=10MB log4j.appender.autoBalancerAppender.MaxBackupIndex=11 log4j.appender.autoBalancerAppender.File=${kafka.logs.dir}/auto-balancer.log @@ -145,3 +148,6 @@ log4j.additivity.kafka.authorizer.logger=false log4j.logger.kafka.autobalancer=INFO, autoBalancerAppender log4j.additivity.kafka.autobalancer=false + +log4j.logger.org.apache.kafka.clients=INFO, kafkaAppender +log4j.additivity.org.apache.kafka.clients=false diff --git a/config/perf-log4j.properties b/config/perf-log4j.properties index 8711f97933..d3f4ac43bf 100644 --- a/config/perf-log4j.properties +++ b/config/perf-log4j.properties @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -log4j.rootLogger=INFO, stdout, perfAppender +log4j.rootLogger=ERROR, stdout, perfAppender log4j.appender.stdout=org.apache.log4j.ConsoleAppender log4j.appender.stdout.layout=org.apache.log4j.PatternLayout @@ -26,7 +26,15 @@ log4j.appender.perfAppender.File=${kafka.logs.dir}/perf.log log4j.appender.perfAppender.layout=org.apache.log4j.PatternLayout log4j.appender.perfAppender.layout.ConversionPattern=%d -%5p [%15.15t] %m (%c#%M:%L)%n -log4j.logger.org.apache.kafka=INFO, perfAppender -log4j.additivity.org.apache.kafka=false +log4j.appender.clientAppender=org.apache.log4j.RollingFileAppender +log4j.appender.clientAppender.MaxFileSize=100MB +log4j.appender.clientAppender.MaxBackupIndex=10 +log4j.appender.clientAppender.File=${kafka.logs.dir}/client.log +log4j.appender.clientAppender.layout=org.apache.log4j.PatternLayout +log4j.appender.clientAppender.layout.ConversionPattern=%d -%5p [%15.15t] %m (%c#%M:%L)%n log4j.logger.org.apache.kafka.tools.automq=INFO, stdout, perfAppender +log4j.additivity.org.apache.kafka.tools.automq=false + +log4j.logger.org.apache.kafka.clients=INFO, clientAppender +log4j.additivity.org.apache.kafka.clients=false diff --git a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/AutoMQIdentityReplicationPolicy.java b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/AutoMQIdentityReplicationPolicy.java index 0fc015172f..1832e68d35 100644 --- a/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/AutoMQIdentityReplicationPolicy.java +++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/AutoMQIdentityReplicationPolicy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.connect.mirror; diff --git a/connect/runtime/README.md b/connect/runtime/README.md new file mode 100644 index 0000000000..2ce7694e5f --- /dev/null +++ b/connect/runtime/README.md @@ -0,0 +1,221 @@ +# Kafka Connect OpenTelemetry Metrics Integration + +## Overview + +This integration allows Kafka Connect to export metrics through the AutoMQ OpenTelemetry module, enabling unified observability across your Kafka ecosystem. + +## Configuration + +### 1. Enable the MetricsReporter + +Add the following to your Kafka Connect configuration file (`connect-distributed.properties` or `connect-standalone.properties`): + +```properties +# Enable OpenTelemetry MetricsReporter +metric.reporters=org.apache.kafka.connect.automq.metrics.OpenTelemetryMetricsReporter + +# OpenTelemetry configuration +opentelemetry.metrics.enabled=true +opentelemetry.metrics.prefix=kafka.connect + +# Optional: Filter metrics +opentelemetry.metrics.include.pattern=.*connector.*|.*task.*|.*worker.* +opentelemetry.metrics.exclude.pattern=.*jmx.*|.*debug.* +``` + +### 2. AutoMQ Telemetry Configuration + +Ensure the AutoMQ telemetry is properly configured. Add these properties to your application configuration: + +```properties +# Telemetry export configuration +automq.telemetry.exporter.uri=prometheus://localhost:9090 +# or for OTLP: automq.telemetry.exporter.uri=otlp://localhost:4317 + +# Service identification +service.name=kafka-connect +service.instance.id=connect-worker-1 + +# Export settings +automq.telemetry.exporter.interval.ms=30000 +automq.telemetry.metric.cardinality.limit=10000 +``` + +## S3 Log Upload + +Kafka Connect bundles the AutoMQ log uploader so that worker logs can be streamed to S3 together with in-cluster cleanup. The uploader uses the connect-leader election mechanism by default and requires no additional configuration. + +### Worker Configuration + +Add the following properties to your worker configuration (ConfigMap, properties file, etc.): + +```properties +# Enable S3 log upload +log.s3.enable=true +log.s3.bucket=0@s3://your-log-bucket?region=us-east-1 + +# Optional overrides (defaults shown) +log.s3.selector.type=connect-leader +# Provide credentials if the bucket URI does not embed them +# log.s3.access.key=... +# log.s3.secret.key=... +``` + +`log.s3.node.id` defaults to a hash of the pod hostname if not provided, ensuring objects are partitioned per worker. + +### Log4j Integration + +`config/connect-log4j.properties` has switched `connectAppender` to `com.automq.log.S3RollingFileAppender` and specifies `org.apache.kafka.connect.automq.log.ConnectS3LogConfigProvider` as the config provider. As long as you enable `log.s3.enable=true` and configure the bucket info in the worker config, log upload will be automatically initialized with the Connect process; if not set or returns `log.s3.enable=false`, the uploader remains disabled. + +## Programmatic Usage + +### 1. Initialize Telemetry Manager + +```java +import com.automq.opentelemetry.AutoMQTelemetryManager; +import java.util.Properties; + +// Initialize AutoMQ telemetry before starting Kafka Connect +Properties telemetryProps = new Properties(); +telemetryProps.setProperty("automq.telemetry.exporter.uri", "prometheus://localhost:9090"); +telemetryProps.setProperty("service.name", "kafka-connect"); +telemetryProps.setProperty("service.instance.id", "worker-1"); + +// Initialize singleton instance +AutoMQTelemetryManager.initializeInstance(telemetryProps); + +// Now start Kafka Connect - it will automatically use the OpenTelemetryMetricsReporter +``` + +### 2. Shutdown + +```java +// When shutting down your application +AutoMQTelemetryManager.shutdownInstance(); +``` + +## Exported Metrics + +The integration automatically converts Kafka Connect metrics to OpenTelemetry format: + +### Metric Naming Convention +- **Format**: `kafka.connect.{group}.{metric_name}` +- **Example**: `kafka.connect.connector.task.batch.size.avg` → `kafka.connect.connector_task_batch_size_avg` + +### Metric Types +- **Counters**: Metrics containing "total", "count", "error", "failure" +- **Gauges**: All other numeric metrics (rates, averages, sizes, etc.) + +### Attributes +Kafka metric tags are converted to OpenTelemetry attributes: +- `connector` → `connector` +- `task` → `task` +- `worker-id` → `worker_id` +- Plus standard attributes: `metric.group`, `service.name`, `service.instance.id` + +## Example Metrics + +Common Kafka Connect metrics that will be exported: + +``` +# Connector metrics +kafka.connect.connector.startup.attempts.total +kafka.connect.connector.startup.success.total +kafka.connect.connector.startup.failure.total + +# Task metrics +kafka.connect.connector.task.batch.size.avg +kafka.connect.connector.task.batch.size.max +kafka.connect.connector.task.offset.commit.avg.time.ms + +# Worker metrics +kafka.connect.worker.connector.count +kafka.connect.worker.task.count +kafka.connect.worker.connector.startup.attempts.total +``` + +## Configuration Options + +### OpenTelemetry MetricsReporter Options + +| Property | Description | Default | Example | +|----------|-------------|---------|---------| +| `opentelemetry.metrics.enabled` | Enable/disable metrics export | `true` | `false` | +| `opentelemetry.metrics.prefix` | Metric name prefix | `kafka.connect` | `my.connect` | +| `opentelemetry.metrics.include.pattern` | Regex for included metrics | All metrics | `.*connector.*` | +| `opentelemetry.metrics.exclude.pattern` | Regex for excluded metrics | None | `.*jmx.*` | + +### AutoMQ Telemetry Options + +| Property | Description | Default | +|----------|-------------|---------| +| `automq.telemetry.exporter.uri` | Exporter endpoint | Empty | +| `automq.telemetry.exporter.interval.ms` | Export interval | `60000` | +| `automq.telemetry.metric.cardinality.limit` | Max metric cardinality | `20000` | + +## Monitoring Examples + +### Prometheus Queries + +```promql +# Connector count by worker +kafka_connect_worker_connector_count + +# Task failure rate +rate(kafka_connect_connector_task_startup_failure_total[5m]) + +# Average batch processing time +kafka_connect_connector_task_batch_size_avg + +# Connector startup success rate +rate(kafka_connect_connector_startup_success_total[5m]) / +rate(kafka_connect_connector_startup_attempts_total[5m]) +``` + +### Grafana Dashboard + +Common panels to create: + +1. **Connector Health**: Count of running/failed connectors +2. **Task Performance**: Batch size, processing time, throughput +3. **Error Rates**: Failed startups, task failures +4. **Resource Usage**: Combined with JVM metrics from AutoMQ telemetry + +## Troubleshooting + +### Common Issues + +1. **Metrics not appearing** + ``` + Check logs for: "AutoMQTelemetryManager is not initialized" + Solution: Ensure AutoMQTelemetryManager.initializeInstance() is called before Connect starts + ``` + +2. **High cardinality warnings** + ``` + Solution: Use include/exclude patterns to filter metrics + ``` + +3. **Missing dependencies** + ``` + Ensure connect-runtime depends on the opentelemetry module + ``` + +### Debug Logging + +Enable debug logging to troubleshoot: + +```properties +log4j.logger.org.apache.kafka.connect.automq=DEBUG +log4j.logger.com.automq.opentelemetry=DEBUG +``` + +## Integration with Existing Monitoring + +This integration works alongside: +- Existing JMX metrics (not replaced) +- Kafka broker metrics via AutoMQ telemetry +- Application-specific metrics +- Third-party monitoring tools + +The OpenTelemetry integration provides a unified export path while preserving existing monitoring setups. diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzAwareClientConfigurator.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzAwareClientConfigurator.java new file mode 100644 index 0000000000..567e93ebe1 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzAwareClientConfigurator.java @@ -0,0 +1,95 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.az; + +import org.apache.kafka.clients.CommonClientConfigs; +import org.apache.kafka.clients.consumer.ConsumerConfig; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import java.util.Optional; + +public final class AzAwareClientConfigurator { + private static final Logger LOGGER = LoggerFactory.getLogger(AzAwareClientConfigurator.class); + + private AzAwareClientConfigurator() { + } + + public enum ClientFamily { + PRODUCER, + CONSUMER, + ADMIN + } + + public static void maybeApplyAz(Map props, ClientFamily family, String roleDescriptor) { + Optional azOpt = AzMetadataProviderHolder.provider().availabilityZoneId(); + LOGGER.info("AZ-aware client.id configuration for role {}: resolved availability zone id '{}'", + roleDescriptor, azOpt.orElse("unknown")); + if (azOpt.isEmpty()) { + LOGGER.info("Skipping AZ-aware client.id configuration for role {} as no availability zone id is available", + roleDescriptor); + return; + } + + String az = azOpt.get(); + + String encodedAz = URLEncoder.encode(az, StandardCharsets.UTF_8); + String automqClientId; + + if (props.containsKey(CommonClientConfigs.CLIENT_ID_CONFIG)) { + Object currentId = props.get(CommonClientConfigs.CLIENT_ID_CONFIG); + if (currentId instanceof String currentIdStr) { + automqClientId = "automq_az=" + encodedAz + "&" + currentIdStr; + } else { + LOGGER.warn("client.id for role {} is not a string ({});", + roleDescriptor, currentId.getClass().getName()); + return; + } + } else { + automqClientId = "automq_az=" + encodedAz; + } + props.put(CommonClientConfigs.CLIENT_ID_CONFIG, automqClientId); + LOGGER.info("Applied AZ-aware client.id for role {} -> {}", roleDescriptor, automqClientId); + + if (family == ClientFamily.CONSUMER) { + LOGGER.info("Applying client.rack configuration for consumer role {} -> {}", roleDescriptor, az); + Object rackValue = props.get(ConsumerConfig.CLIENT_RACK_CONFIG); + if (rackValue == null || String.valueOf(rackValue).isBlank()) { + props.put(ConsumerConfig.CLIENT_RACK_CONFIG, az); + } + } + } + + public static void maybeApplyProducerAz(Map props, String roleDescriptor) { + maybeApplyAz(props, ClientFamily.PRODUCER, roleDescriptor); + } + + public static void maybeApplyConsumerAz(Map props, String roleDescriptor) { + maybeApplyAz(props, ClientFamily.CONSUMER, roleDescriptor); + } + + public static void maybeApplyAdminAz(Map props, String roleDescriptor) { + maybeApplyAz(props, ClientFamily.ADMIN, roleDescriptor); + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzMetadataProvider.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzMetadataProvider.java new file mode 100644 index 0000000000..295d6e9f22 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzMetadataProvider.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.az; + +import java.util.Map; +import java.util.Optional; + +/** + * Pluggable provider for availability-zone metadata used to tune Kafka client configurations. + */ +public interface AzMetadataProvider { + + /** + * Configure the provider with the worker properties. Implementations may cache values extracted from the + * configuration map. This method is invoked exactly once during worker bootstrap. + */ + default void configure(Map workerProps) { + // no-op + } + + /** + * @return the availability-zone identifier for the current node, if known. + */ + default Optional availabilityZoneId() { + return Optional.empty(); + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzMetadataProviderHolder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzMetadataProviderHolder.java new file mode 100644 index 0000000000..861aa4dbd5 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/az/AzMetadataProviderHolder.java @@ -0,0 +1,64 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.az; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.ServiceLoader; + +public final class AzMetadataProviderHolder { + private static final Logger LOGGER = LoggerFactory.getLogger(AzMetadataProviderHolder.class); + private static final AzMetadataProvider DEFAULT_PROVIDER = new AzMetadataProvider() { }; + + private static volatile AzMetadataProvider provider = DEFAULT_PROVIDER; + + private AzMetadataProviderHolder() { + } + + public static void initialize(Map workerProps) { + AzMetadataProvider selected = DEFAULT_PROVIDER; + try { + ServiceLoader loader = ServiceLoader.load(AzMetadataProvider.class); + for (AzMetadataProvider candidate : loader) { + try { + candidate.configure(workerProps); + selected = candidate; + LOGGER.info("Loaded AZ metadata provider: {}", candidate.getClass().getName()); + break; + } catch (Exception e) { + LOGGER.warn("Failed to initialize AZ metadata provider: {}", candidate.getClass().getName(), e); + } + } + } catch (Throwable t) { + LOGGER.warn("Failed to load AZ metadata providers", t); + } + provider = selected; + } + + public static AzMetadataProvider provider() { + return provider; + } + + public static void setProviderForTest(AzMetadataProvider newProvider) { + provider = newProvider != null ? newProvider : DEFAULT_PROVIDER; + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectLogUploader.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectLogUploader.java new file mode 100644 index 0000000000..166859b4fb --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectLogUploader.java @@ -0,0 +1,56 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.log; + +import com.automq.log.S3RollingFileAppender; +import com.automq.log.uploader.S3LogConfig; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.Properties; + +/** + * Initializes the AutoMQ S3 log uploader for Kafka Connect. + */ +public final class ConnectLogUploader { + private static Logger getLogger() { + return LoggerFactory.getLogger(ConnectLogUploader.class); + } + + private ConnectLogUploader() { + } + + public static void initialize(Map workerProps) { + Properties props = new Properties(); + if (workerProps != null) { + workerProps.forEach((k, v) -> { + if (k != null && v != null) { + props.put(k, v); + } + }); + } + ConnectS3LogConfigProvider.initialize(props); + S3LogConfig s3LogConfig = new ConnectS3LogConfigProvider().get(); + S3RollingFileAppender.setup(s3LogConfig); + getLogger().info("Initialized Connect S3 log uploader context"); + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectS3LogConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectS3LogConfig.java new file mode 100644 index 0000000000..3627cd3bed --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectS3LogConfig.java @@ -0,0 +1,95 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.log; + +import org.apache.kafka.connect.automq.runtime.LeaderNodeSelector; +import org.apache.kafka.connect.automq.runtime.RuntimeLeaderSelectorProvider; + +import com.automq.log.uploader.S3LogConfig; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; + +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ConnectS3LogConfig implements S3LogConfig { + private static final Logger LOGGER = LoggerFactory.getLogger(ConnectS3LogConfig.class); + + private final boolean enable; + private final String clusterId; + private final int nodeId; + private final String bucketURI; + private ObjectStorage objectStorage; + private LeaderNodeSelector leaderNodeSelector; + + + public ConnectS3LogConfig(boolean enable, String clusterId, int nodeId, String bucketURI) { + this.enable = enable; + this.clusterId = clusterId; + this.nodeId = nodeId; + this.bucketURI = bucketURI; + } + + @Override + public boolean isEnabled() { + return this.enable; + } + + @Override + public String clusterId() { + return this.clusterId; + } + + @Override + public int nodeId() { + return this.nodeId; + } + + @Override + public synchronized ObjectStorage objectStorage() { + if (this.objectStorage != null) { + return this.objectStorage; + } + if (StringUtils.isBlank(bucketURI)) { + LOGGER.error("Mandatory log config bucketURI is not set."); + return null; + } + + String normalizedBucket = bucketURI.trim(); + BucketURI logBucket = BucketURI.parse(normalizedBucket); + this.objectStorage = ObjectStorageFactory.instance().builder(logBucket).threadPrefix("s3-log-uploader").build(); + return this.objectStorage; + } + + @Override + public boolean isLeader() { + LeaderNodeSelector selector = leaderSelector(); + return selector != null && selector.isLeader(); + } + + public LeaderNodeSelector leaderSelector() { + if (leaderNodeSelector == null) { + this.leaderNodeSelector = new RuntimeLeaderSelectorProvider().createSelector(); + } + return leaderNodeSelector; + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectS3LogConfigProvider.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectS3LogConfigProvider.java new file mode 100644 index 0000000000..acb7c56ec8 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/ConnectS3LogConfigProvider.java @@ -0,0 +1,112 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.log; + +import com.automq.log.uploader.S3LogConfig; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.InetAddress; +import java.util.Map; +import java.util.Properties; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Provides S3 log uploader configuration for Kafka Connect workers. + */ +public class ConnectS3LogConfigProvider { + private static Logger getLogger() { + return LoggerFactory.getLogger(ConnectS3LogConfigProvider.class); + } + private static final AtomicReference CONFIG = new AtomicReference<>(); + private static final long WAIT_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(10); + private static final CountDownLatch INIT = new CountDownLatch(1); + + public static void initialize(Properties workerProps) { + try { + if (workerProps == null) { + CONFIG.set(null); + return; + } + Properties copy = new Properties(); + for (Map.Entry entry : workerProps.entrySet()) { + if (entry.getKey() != null && entry.getValue() != null) { + copy.put(entry.getKey(), entry.getValue()); + } + } + CONFIG.set(copy); + } finally { + INIT.countDown(); + } + getLogger().info("Initializing ConnectS3LogConfigProvider"); + } + + public S3LogConfig get() { + + try { + if (!INIT.await(WAIT_TIMEOUT_MS, TimeUnit.MILLISECONDS)) { + getLogger().warn("S3 log uploader config not initialized within timeout; uploader disabled."); + } + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + getLogger().warn("Interrupted while waiting for S3 log uploader config; uploader disabled."); + return null; + } + + Properties source = CONFIG.get(); + if (source == null) { + getLogger().warn("S3 log upload configuration was not provided; uploader disabled."); + return null; + } + + String bucketURI = source.getProperty(LogConfigConstants.LOG_S3_BUCKET_KEY); + String clusterId = source.getProperty(LogConfigConstants.LOG_S3_CLUSTER_ID_KEY); + String nodeIdStr = resolveNodeId(source); + boolean enable = Boolean.parseBoolean(source.getProperty(LogConfigConstants.LOG_S3_ENABLE_KEY, "false")); + return new ConnectS3LogConfig(enable, clusterId, Integer.parseInt(nodeIdStr), bucketURI); + } + + private String resolveNodeId(Properties workerProps) { + String fromConfig = workerProps.getProperty(LogConfigConstants.LOG_S3_NODE_ID_KEY); + if (!isBlank(fromConfig)) { + return fromConfig.trim(); + } + String env = System.getenv("CONNECT_NODE_ID"); + if (!isBlank(env)) { + return env.trim(); + } + String host = workerProps.getProperty("automq.log.s3.node.hostname"); + if (isBlank(host)) { + try { + host = InetAddress.getLocalHost().getHostName(); + } catch (Exception e) { + host = System.getenv().getOrDefault("HOSTNAME", "0"); + } + } + return Integer.toString(host.hashCode() & Integer.MAX_VALUE); + } + + private boolean isBlank(String value) { + return value == null || value.trim().isEmpty(); + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/LogConfigConstants.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/LogConfigConstants.java new file mode 100644 index 0000000000..191400f457 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/log/LogConfigConstants.java @@ -0,0 +1,30 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.log; + +public class LogConfigConstants { + public static final String LOG_S3_ENABLE_KEY = "log.s3.enable"; + + public static final String LOG_S3_BUCKET_KEY = "log.s3.bucket"; + + public static final String LOG_S3_CLUSTER_ID_KEY = "log.s3.cluster.id"; + + public static final String LOG_S3_NODE_ID_KEY = "log.s3.node.id"; +} \ No newline at end of file diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/ConnectMetricsExportConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/ConnectMetricsExportConfig.java new file mode 100644 index 0000000000..86d6bdeb14 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/ConnectMetricsExportConfig.java @@ -0,0 +1,77 @@ +package org.apache.kafka.connect.automq.metrics; + +import org.apache.kafka.connect.automq.runtime.LeaderNodeSelector; +import org.apache.kafka.connect.automq.runtime.RuntimeLeaderSelectorProvider; + +import com.automq.opentelemetry.exporter.MetricsExportConfig; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; + +import org.apache.commons.lang3.tuple.Pair; + +import java.util.List; + +public class ConnectMetricsExportConfig implements MetricsExportConfig { + + private final BucketURI metricsBucket; + private final String clusterId; + private final int nodeId; + private final int intervalMs; + private final List> baseLabels; + private ObjectStorage objectStorage; + private LeaderNodeSelector leaderNodeSelector; + + + public ConnectMetricsExportConfig(String clusterId, int nodeId, BucketURI metricsBucket, List> baseLabels, int intervalMs) { + this.clusterId = clusterId; + this.nodeId = nodeId; + this.metricsBucket = metricsBucket; + this.baseLabels = baseLabels; + this.intervalMs = intervalMs; + } + + @Override + public String clusterId() { + return this.clusterId; + } + + @Override + public boolean isLeader() { + LeaderNodeSelector selector = leaderSelector(); + return selector != null && selector.isLeader(); + } + + public LeaderNodeSelector leaderSelector() { + if (leaderNodeSelector == null) { + this.leaderNodeSelector = new RuntimeLeaderSelectorProvider().createSelector(); + } + return leaderNodeSelector; + } + + @Override + public int nodeId() { + return this.nodeId; + } + + @Override + public ObjectStorage objectStorage() { + if (metricsBucket == null) { + return null; + } + if (this.objectStorage == null) { + this.objectStorage = ObjectStorageFactory.instance().builder(metricsBucket).threadPrefix("s3-metric").build(); + } + return this.objectStorage; + } + + @Override + public List> baseLabels() { + return this.baseLabels; + } + + @Override + public int intervalMs() { + return this.intervalMs; + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/MetricsConfigConstants.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/MetricsConfigConstants.java new file mode 100644 index 0000000000..4e34467989 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/MetricsConfigConstants.java @@ -0,0 +1,30 @@ +package org.apache.kafka.connect.automq.metrics; + +public class MetricsConfigConstants { + public static final String SERVICE_NAME_KEY = "service.name"; + public static final String SERVICE_INSTANCE_ID_KEY = "service.instance.id"; + public static final String S3_CLIENT_ID_KEY = "automq.telemetry.s3.cluster.id"; + /** + * The URI for configuring metrics exporters. e.g. prometheus://localhost:9090, otlp://localhost:4317 + */ + public static final String EXPORTER_URI_KEY = "automq.telemetry.exporter.uri"; + /** + * The export interval in milliseconds. + */ + public static final String EXPORTER_INTERVAL_MS_KEY = "automq.telemetry.exporter.interval.ms"; + /** + * The cardinality limit for any single metric. + */ + public static final String METRIC_CARDINALITY_LIMIT_KEY = "automq.telemetry.metric.cardinality.limit"; + public static final int DEFAULT_METRIC_CARDINALITY_LIMIT = 20000; + + public static final String TELEMETRY_METRICS_BASE_LABELS_CONFIG = "automq.telemetry.metrics.base.labels"; + public static final String TELEMETRY_METRICS_BASE_LABELS_DOC = "The base labels that will be added to all metrics. The format is key1=value1,key2=value2."; + + public static final String S3_BUCKET = "automq.telemetry.s3.bucket"; + public static final String S3_BUCKETS_DOC = "The buckets url with format 0@s3://$bucket?region=$region. \n" + + "the full url format for s3 is 0@s3://$bucket?region=$region[&endpoint=$endpoint][&pathStyle=$enablePathStyle][&authType=$authType][&accessKey=$accessKey][&secretKey=$secretKey][&checksumAlgorithm=$checksumAlgorithm]" + + "- pathStyle: true|false. The object storage access path style. When using MinIO, it should be set to true.\n" + + "- authType: instance|static. When set to instance, it will use instance profile to auth. When set to static, it will get accessKey and secretKey from the url or from system environment KAFKA_S3_ACCESS_KEY/KAFKA_S3_SECRET_KEY."; + +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/OpenTelemetryMetricsReporter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/OpenTelemetryMetricsReporter.java new file mode 100644 index 0000000000..1f43471f57 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/metrics/OpenTelemetryMetricsReporter.java @@ -0,0 +1,822 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.metrics; + +import org.apache.kafka.common.MetricName; +import org.apache.kafka.common.metrics.KafkaMetric; +import org.apache.kafka.common.metrics.MetricsReporter; + +import com.automq.opentelemetry.AutoMQTelemetryManager; +import com.automq.stream.s3.operator.BucketURI; + +import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.tuple.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Properties; +import java.util.concurrent.ConcurrentHashMap; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.common.AttributesBuilder; +import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.metrics.ObservableDoubleCounter; +import io.opentelemetry.api.metrics.ObservableDoubleGauge; +import io.opentelemetry.api.metrics.ObservableLongCounter; + +/** + * A MetricsReporter implementation that bridges Kafka Connect metrics to OpenTelemetry. + * + *

This reporter integrates with the AutoMQ OpenTelemetry module to export Kafka Connect + * metrics through various exporters (Prometheus, OTLP, etc.). It automatically converts + * Kafka metrics to OpenTelemetry instruments based on metric types and provides proper + * labeling and naming conventions. + * + *

Key features: + *

    + *
  • Automatic metric type detection and conversion
    • + *
  • Support for gauges and counters using async observable instruments
    • + *
  • Proper attribute mapping from Kafka metric tags
    • + *
  • Integration with AutoMQ telemetry infrastructure
    • + *
  • Configurable metric filtering
    • + *
  • Real-time metric value updates through callbacks
    • + *
+ * + *

Configuration options: + *

    + *
  • {@code opentelemetry.metrics.enabled} - Enable/disable OpenTelemetry metrics (default: true)
    • + *
  • {@code opentelemetry.metrics.prefix} - Prefix for metric names (default: "kafka.connect")
    • + *
  • {@code opentelemetry.metrics.include.pattern} - Regex pattern for included metrics
    • + *
  • {@code opentelemetry.metrics.exclude.pattern} - Regex pattern for excluded metrics
    • + *
+ */ +public class OpenTelemetryMetricsReporter implements MetricsReporter { + private static final Logger LOGGER = LoggerFactory.getLogger(OpenTelemetryMetricsReporter.class); + + private static final String ENABLED_CONFIG = "opentelemetry.metrics.enabled"; + private static final String PREFIX_CONFIG = "opentelemetry.metrics.prefix"; + private static final String INCLUDE_PATTERN_CONFIG = "opentelemetry.metrics.include.pattern"; + private static final String EXCLUDE_PATTERN_CONFIG = "opentelemetry.metrics.exclude.pattern"; + + private static final String DEFAULT_PREFIX = "kafka"; + + private boolean enabled = true; + private String metricPrefix = DEFAULT_PREFIX; + private String includePattern = null; + private String excludePattern = null; + + private Meter meter; + private final Map observableHandles = new ConcurrentHashMap<>(); + private final Map registeredMetrics = new ConcurrentHashMap<>(); + + public static void initializeTelemetry(Properties props) { + String exportURIStr = props.getProperty(MetricsConfigConstants.EXPORTER_URI_KEY); + String serviceName = props.getProperty(MetricsConfigConstants.SERVICE_NAME_KEY, "connect-default"); + String instanceId = props.getProperty(MetricsConfigConstants.SERVICE_INSTANCE_ID_KEY, "0"); + String clusterId = props.getProperty(MetricsConfigConstants.S3_CLIENT_ID_KEY, "cluster-default"); + int intervalMs = Integer.parseInt(props.getProperty(MetricsConfigConstants.EXPORTER_INTERVAL_MS_KEY, "60000")); + BucketURI metricsBucket = getMetricsBucket(props); + List> baseLabels = getBaseLabels(props); + + AutoMQTelemetryManager.initializeInstance(exportURIStr, serviceName, instanceId, new ConnectMetricsExportConfig(clusterId, Integer.parseInt(instanceId), metricsBucket, baseLabels, intervalMs)); + LOGGER.info("OpenTelemetryMetricsReporter initialized"); + } + + private static BucketURI getMetricsBucket(Properties props) { + String metricsBucket = props.getProperty(MetricsConfigConstants.S3_BUCKET, ""); + if (StringUtils.isNotBlank(metricsBucket)) { + List bucketList = BucketURI.parseBuckets(metricsBucket); + if (!bucketList.isEmpty()) { + return bucketList.get(0); + } + } + return null; + } + + private static List> getBaseLabels(Properties props) { + // This part is hard to abstract without a clear config pattern. + // Assuming for now it's empty. The caller can extend this class + // or the manager can have a method to add more labels. + String baseLabels = props.getProperty(MetricsConfigConstants.TELEMETRY_METRICS_BASE_LABELS_CONFIG); + if (StringUtils.isBlank(baseLabels)) { + return Collections.emptyList(); + } + List> labels = new ArrayList<>(); + for (String label : baseLabels.split(",")) { + String[] kv = label.split("="); + if (kv.length != 2) { + continue; + } + labels.add(Pair.of(kv[0], kv[1])); + } + return labels; + } + + @Override + public void configure(Map configs) { + // Parse configuration + Object enabledObj = configs.get(ENABLED_CONFIG); + if (enabledObj != null) { + enabled = Boolean.parseBoolean(enabledObj.toString()); + } + + Object prefixObj = configs.get(PREFIX_CONFIG); + if (prefixObj != null) { + metricPrefix = prefixObj.toString(); + } + + Object includeObj = configs.get(INCLUDE_PATTERN_CONFIG); + if (includeObj != null) { + includePattern = includeObj.toString(); + } + + Object excludeObj = configs.get(EXCLUDE_PATTERN_CONFIG); + if (excludeObj != null) { + excludePattern = excludeObj.toString(); + } + + LOGGER.info("OpenTelemetryMetricsReporter configured - enabled: {}, prefix: {}, include: {}, exclude: {}", + enabled, metricPrefix, includePattern, excludePattern); + } + + @Override + public void init(List metrics) { + if (!enabled) { + LOGGER.info("OpenTelemetryMetricsReporter is disabled"); + return; + } + + try { + // Get the OpenTelemetry meter from AutoMQTelemetryManager + // This assumes the telemetry manager is already initialized + meter = AutoMQTelemetryManager.getInstance().getMeter(); + if (meter == null) { + LOGGER.warn("AutoMQTelemetryManager is not initialized, OpenTelemetry metrics will not be available"); + enabled = false; + return; + } + + // Register initial metrics + for (KafkaMetric metric : metrics) { + registerMetric(metric); + } + + LOGGER.info("OpenTelemetryMetricsReporter initialized with {} metrics", metrics.size()); + } catch (Exception e) { + LOGGER.error("Failed to initialize OpenTelemetryMetricsReporter", e); + enabled = false; + } + } + + @Override + public void metricChange(KafkaMetric metric) { + if (!enabled || meter == null) { + return; + } + + try { + registerMetric(metric); + } catch (Exception e) { + LOGGER.warn("Failed to register metric change for {}", metric.metricName(), e); + } + } + + @Override + public void metricRemoval(KafkaMetric metric) { + if (!enabled) { + return; + } + + try { + String metricKey = buildMetricKey(metric.metricName()); + closeHandle(metricKey); + registeredMetrics.remove(metricKey); + LOGGER.debug("Removed metric: {}", metricKey); + } catch (Exception e) { + LOGGER.warn("Failed to remove metric {}", metric.metricName(), e); + } + } + + @Override + public void close() { + if (enabled) { + // Close all observable handles to prevent memory leaks + observableHandles.values().forEach(handle -> { + try { + handle.close(); + } catch (Exception e) { + LOGGER.debug("Error closing observable handle", e); + } + }); + observableHandles.clear(); + registeredMetrics.clear(); + } + LOGGER.info("OpenTelemetryMetricsReporter closed"); + } + + private void registerMetric(KafkaMetric metric) { + LOGGER.debug("OpenTelemetryMetricsReporter registering metric {}", metric.metricName()); + MetricName metricName = metric.metricName(); + String metricKey = buildMetricKey(metricName); + + // Apply filtering + if (!shouldIncludeMetric(metricKey)) { + return; + } + + // Check if metric value is numeric at registration time + Object testValue = safeMetricValue(metric); + if (!(testValue instanceof Number)) { + LOGGER.debug("Skipping non-numeric metric: {}", metricKey); + return; + } + + Attributes attributes = buildAttributes(metricName); + + // Close existing handle if present (for metric updates) + closeHandle(metricKey); + + // Register the metric for future access + registeredMetrics.put(metricKey, metric); + + // Determine metric type and register accordingly + if (isCounterMetric(metricName)) { + registerAsyncCounter(metricKey, metricName, metric, attributes, (Number) testValue); + } else { + registerAsyncGauge(metricKey, metricName, metric, attributes); + } + } + + private void registerAsyncGauge(String metricKey, MetricName metricName, KafkaMetric metric, Attributes attributes) { + try { + String description = buildDescription(metricName); + String unit = determineUnit(metricName); + + ObservableDoubleGauge gauge = meter.gaugeBuilder(metricKey) + .setDescription(description) + .setUnit(unit) + .buildWithCallback(measurement -> { + Number value = (Number) safeMetricValue(metric); + if (value != null) { + measurement.record(value.doubleValue(), attributes); + } + }); + + observableHandles.put(metricKey, gauge); + LOGGER.debug("Registered async gauge: {}", metricKey); + } catch (Exception e) { + LOGGER.warn("Failed to register async gauge for {}", metricKey, e); + } + } + + private void registerAsyncCounter(String metricKey, MetricName metricName, KafkaMetric metric, + Attributes attributes, Number initialValue) { + try { + String description = buildDescription(metricName); + String unit = determineUnit(metricName); + + // Use appropriate counter type based on initial value type + if (initialValue instanceof Long || initialValue instanceof Integer) { + ObservableLongCounter counter = meter.counterBuilder(metricKey) + .setDescription(description) + .setUnit(unit) + .buildWithCallback(measurement -> { + Number value = (Number) safeMetricValue(metric); + if (value != null) { + long longValue = value.longValue(); + if (longValue >= 0) { + measurement.record(longValue, attributes); + } + } + }); + observableHandles.put(metricKey, counter); + } else { + ObservableDoubleCounter counter = meter.counterBuilder(metricKey) + .ofDoubles() + .setDescription(description) + .setUnit(unit) + .buildWithCallback(measurement -> { + Number value = (Number) safeMetricValue(metric); + if (value != null) { + double doubleValue = value.doubleValue(); + if (doubleValue >= 0) { + measurement.record(doubleValue, attributes); + } + } + }); + observableHandles.put(metricKey, counter); + } + + LOGGER.debug("Registered async counter: {}", metricKey); + } catch (Exception e) { + LOGGER.warn("Failed to register async counter for {}", metricKey, e); + } + } + + private Object safeMetricValue(KafkaMetric metric) { + try { + return metric.metricValue(); + } catch (Exception e) { + LOGGER.debug("Failed to read metric value for {}", metric.metricName(), e); + return null; + } + } + + private void closeHandle(String metricKey) { + AutoCloseable handle = observableHandles.remove(metricKey); + if (handle != null) { + try { + handle.close(); + } catch (Exception e) { + LOGGER.debug("Error closing handle for {}", metricKey, e); + } + } + } + + private String buildMetricKey(MetricName metricName) { + StringBuilder sb = new StringBuilder(metricPrefix); + sb.append("."); + + // Add group if present + if (metricName.group() != null && !metricName.group().isEmpty()) { + sb.append(metricName.group().replace("-", "_").toLowerCase(Locale.ROOT)); + sb.append("."); + } + + // Add name + sb.append(metricName.name().replace("-", "_").toLowerCase(Locale.ROOT)); + + return sb.toString(); + } + + private Attributes buildAttributes(MetricName metricName) { + AttributesBuilder builder = Attributes.builder(); + + // Add metric tags as attributes + Map tags = metricName.tags(); + if (tags != null) { + for (Map.Entry entry : tags.entrySet()) { + String key = entry.getKey(); + String value = entry.getValue(); + if (key != null && value != null) { + builder.put(sanitizeAttributeKey(key), value); + } + } + } + + // Add standard attributes + if (metricName.group() != null) { + builder.put("metric.group", metricName.group()); + } + + return builder.build(); + } + + private String sanitizeAttributeKey(String key) { + return key.replace("-", "_").replace(".", "_").toLowerCase(Locale.ROOT); + } + + private String buildDescription(MetricName metricName) { + StringBuilder description = new StringBuilder(); + description.append("Kafka Connect metric: "); + + if (metricName.group() != null) { + description.append(metricName.group()).append(" - "); + } + + description.append(metricName.name()); + + return description.toString(); + } + + private String determineUnit(MetricName metricName) { + String name = metricName.name().toLowerCase(Locale.ROOT); + String group = metricName.group() != null ? metricName.group().toLowerCase(Locale.ROOT) : ""; + + if (isKafkaConnectMetric(group)) { + return determineConnectMetricUnit(name); + } + + if (isTimeMetric(name)) { + return determineTimeUnit(name); + } + + if (isBytesMetric(name)) { + return determineBytesUnit(name); + } + + if (isRateMetric(name)) { + return "1/s"; + } + + if (isRatioOrPercentageMetric(name)) { + return "1"; + } + + if (isCountMetric(name)) { + return "1"; + } + + return "1"; + } + + private boolean isCounterMetric(MetricName metricName) { + String name = metricName.name().toLowerCase(Locale.ROOT); + String group = metricName.group() != null ? metricName.group().toLowerCase(Locale.ROOT) : ""; + + if (isKafkaConnectMetric(group)) { + return isConnectCounterMetric(name); + } + + if (isGaugeMetric(name)) { + return false; + } + + return hasCounterKeywords(name); + } + + private boolean isGaugeMetric(String name) { + return hasRateOrAvgKeywords(name) || hasRatioOrPercentKeywords(name) || + hasMinMaxOrCurrentKeywords(name) || hasActiveOrSizeKeywords(name) || + hasTimeButNotTotal(name); + } + + private boolean hasRateOrAvgKeywords(String name) { + return name.contains("rate") || name.contains("avg") || name.contains("mean"); + } + + private boolean hasRatioOrPercentKeywords(String name) { + return name.contains("ratio") || name.contains("percent") || name.contains("pct"); + } + + private boolean hasMinMaxOrCurrentKeywords(String name) { + return name.contains("max") || name.contains("min") || name.contains("current"); + } + + private boolean hasActiveOrSizeKeywords(String name) { + return name.contains("active") || name.contains("lag") || name.contains("size"); + } + + private boolean hasTimeButNotTotal(String name) { + return name.contains("time") && !name.contains("total"); + } + + private boolean hasCounterKeywords(String name) { + String[] parts = name.split("[._-]"); + for (String part : parts) { + if (isCounterKeyword(part)) { + return true; + } + } + return false; + } + + private boolean isCounterKeyword(String part) { + return isBasicCounterKeyword(part) || isAdvancedCounterKeyword(part); + } + + private boolean isBasicCounterKeyword(String part) { + return "total".equals(part) || "count".equals(part) || "sum".equals(part) || + "attempts".equals(part); + } + + private boolean isAdvancedCounterKeyword(String part) { + return "success".equals(part) || "failure".equals(part) || + "errors".equals(part) || "retries".equals(part) || "skipped".equals(part); + } + + private boolean isConnectCounterMetric(String name) { + if (hasTotalBasedCounters(name)) { + return true; + } + + if (hasRecordCounters(name)) { + return true; + } + + if (hasActiveCountMetrics(name)) { + return false; + } + + return false; + } + + private boolean hasTotalBasedCounters(String name) { + return hasBasicTotalCounters(name) || hasSuccessFailureCounters(name) || + hasErrorRetryCounters(name) || hasRequestCompletionCounters(name); + } + + private boolean hasBasicTotalCounters(String name) { + return name.contains("total") || name.contains("attempts"); + } + + private boolean hasSuccessFailureCounters(String name) { + return (name.contains("success") && name.contains("total")) || + (name.contains("failure") && name.contains("total")); + } + + private boolean hasErrorRetryCounters(String name) { + return name.contains("errors") || name.contains("retries") || name.contains("skipped"); + } + + private boolean hasRequestCompletionCounters(String name) { + return name.contains("requests") || name.contains("completions"); + } + + private boolean hasRecordCounters(String name) { + return hasRecordKeyword(name) && hasTotalOperation(name); + } + + private boolean hasRecordKeyword(String name) { + return name.contains("record") || name.contains("records"); + } + + private boolean hasTotalOperation(String name) { + return hasPollWriteTotal(name) || hasReadSendTotal(name); + } + + private boolean hasPollWriteTotal(String name) { + return name.contains("poll-total") || name.contains("write-total"); + } + + private boolean hasReadSendTotal(String name) { + return name.contains("read-total") || name.contains("send-total"); + } + + private boolean hasActiveCountMetrics(String name) { + return hasCountMetrics(name) || hasSequenceMetrics(name); + } + + private boolean hasCountMetrics(String name) { + return hasActiveTaskCount(name) || hasConnectorCount(name) || hasStatusCount(name); + } + + private boolean hasActiveTaskCount(String name) { + return name.contains("active-count") || name.contains("partition-count") || + name.contains("task-count"); + } + + private boolean hasConnectorCount(String name) { + return name.contains("connector-count") || name.contains("running-count"); + } + + private boolean hasStatusCount(String name) { + return name.contains("paused-count") || name.contains("failed-count"); + } + + private boolean hasSequenceMetrics(String name) { + return name.contains("seq-no") || name.contains("seq-num"); + } + + private boolean isKafkaConnectMetric(String group) { + return group.contains("connector") || group.contains("task") || + group.contains("connect") || group.contains("worker"); + } + + private String determineConnectMetricUnit(String name) { + String timeUnit = getTimeUnit(name); + if (timeUnit != null) { + return timeUnit; + } + + String countUnit = getCountUnit(name); + if (countUnit != null) { + return countUnit; + } + + String specialUnit = getSpecialUnit(name); + if (specialUnit != null) { + return specialUnit; + } + + return "1"; + } + + private String getTimeUnit(String name) { + if (isTimeBasedMetric(name)) { + return "ms"; + } + if (isTimestampMetric(name)) { + return "ms"; + } + if (isTimeSinceMetric(name)) { + return "ms"; + } + return null; + } + + private String getCountUnit(String name) { + if (isSequenceOrCountMetric(name)) { + return "1"; + } + if (isLagMetric(name)) { + return "1"; + } + if (isTotalOrCounterMetric(name)) { + return "1"; + } + return null; + } + + private String getSpecialUnit(String name) { + if (isStatusOrMetadataMetric(name)) { + return "1"; + } + if (isConnectRateMetric(name)) { + return "1/s"; + } + if (isRatioMetric(name)) { + return "1"; + } + return null; + } + + private boolean isTimeBasedMetric(String name) { + return hasTimeMs(name) || hasCommitBatchTime(name); + } + + private boolean hasTimeMs(String name) { + return name.endsWith("-time-ms") || name.endsWith("-avg-time-ms") || + name.endsWith("-max-time-ms"); + } + + private boolean hasCommitBatchTime(String name) { + return name.contains("commit-time") || name.contains("batch-time") || + name.contains("rebalance-time"); + } + + private boolean isSequenceOrCountMetric(String name) { + return hasSequenceNumbers(name) || hasCountSuffix(name); + } + + private boolean hasSequenceNumbers(String name) { + return name.contains("seq-no") || name.contains("seq-num"); + } + + private boolean hasCountSuffix(String name) { + return name.endsWith("-count") || name.contains("task-count") || + name.contains("partition-count"); + } + + private boolean isLagMetric(String name) { + return name.contains("lag"); + } + + private boolean isStatusOrMetadataMetric(String name) { + return isStatusMetric(name) || hasProtocolLeaderMetrics(name) || + hasConnectorMetrics(name); + } + + private boolean isStatusMetric(String name) { + return "status".equals(name) || name.contains("protocol"); + } + + private boolean hasProtocolLeaderMetrics(String name) { + return name.contains("leader-name"); + } + + private boolean hasConnectorMetrics(String name) { + return name.contains("connector-type") || name.contains("connector-class") || + name.contains("connector-version"); + } + + private boolean isRatioMetric(String name) { + return name.contains("ratio") || name.contains("percentage"); + } + + private boolean isTotalOrCounterMetric(String name) { + return hasTotalSum(name) || hasAttempts(name) || hasSuccessFailure(name) || + hasErrorsRetries(name); + } + + private boolean hasTotalSum(String name) { + return name.contains("total") || name.contains("sum"); + } + + private boolean hasAttempts(String name) { + return name.contains("attempts"); + } + + private boolean hasSuccessFailure(String name) { + return name.contains("success") || name.contains("failure"); + } + + private boolean hasErrorsRetries(String name) { + return name.contains("errors") || name.contains("retries") || name.contains("skipped"); + } + + private boolean isTimestampMetric(String name) { + return name.contains("timestamp") || name.contains("epoch"); + } + + private boolean isConnectRateMetric(String name) { + return name.contains("rate") && !name.contains("ratio"); + } + + private boolean isTimeSinceMetric(String name) { + return name.contains("time-since-last") || name.contains("since-last"); + } + + private boolean isTimeMetric(String name) { + return hasTimeKeywords(name) && !hasTimeExclusions(name); + } + + private boolean hasTimeKeywords(String name) { + return name.contains("time") || name.contains("latency") || + name.contains("duration"); + } + + private boolean hasTimeExclusions(String name) { + return name.contains("ratio") || name.contains("rate") || + name.contains("count") || name.contains("since-last"); + } + + private String determineTimeUnit(String name) { + if (name.contains("ms") || name.contains("millisecond")) { + return "ms"; + } else if (name.contains("us") || name.contains("microsecond")) { + return "us"; + } else if (name.contains("ns") || name.contains("nanosecond")) { + return "ns"; + } else if (name.contains("s") && !name.contains("ms")) { + return "s"; + } else { + return "ms"; + } + } + + private boolean isBytesMetric(String name) { + return name.contains("byte") || name.contains("bytes") || + name.contains("size") && !name.contains("batch-size"); + } + + private String determineBytesUnit(String name) { + boolean isRate = name.contains("rate") || name.contains("per-sec") || + name.contains("persec") || name.contains("/s"); + return isRate ? "By/s" : "By"; + } + + private boolean isRateMetric(String name) { + return hasRateKeywords(name) && !hasExcludedKeywords(name); + } + + private boolean hasRateKeywords(String name) { + return name.contains("rate") || name.contains("per-sec") || + name.contains("persec") || name.contains("/s"); + } + + private boolean hasExcludedKeywords(String name) { + return name.contains("byte") || name.contains("ratio"); + } + + private boolean isRatioOrPercentageMetric(String name) { + return hasPercentKeywords(name) || hasRatioKeywords(name); + } + + private boolean hasPercentKeywords(String name) { + return name.contains("percent") || name.contains("pct"); + } + + private boolean hasRatioKeywords(String name) { + return name.contains("ratio"); + } + + private boolean isCountMetric(String name) { + return name.contains("count") || name.contains("total") || + name.contains("sum") || name.endsWith("-num"); + } + + private boolean shouldIncludeMetric(String metricKey) { + if (excludePattern != null && metricKey.matches(excludePattern)) { + return false; + } + + if (includePattern != null) { + return metricKey.matches(includePattern); + } + + return true; + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/LeaderNodeSelector.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/LeaderNodeSelector.java new file mode 100644 index 0000000000..111d50f511 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/LeaderNodeSelector.java @@ -0,0 +1,34 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.runtime; + +/** + * An interface for determining which node should be responsible for clean metrics. + * This abstraction allows different implementations of clean node selection strategies. + */ +public interface LeaderNodeSelector { + + /** + * Determines if the current node should be responsible for clean metrics. + * + * @return true if the current node should clean metrics, false otherwise. + */ + boolean isLeader(); +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/LeaderNodeSelectorProvider.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/LeaderNodeSelectorProvider.java new file mode 100644 index 0000000000..b9bac44957 --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/LeaderNodeSelectorProvider.java @@ -0,0 +1,36 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.runtime; + +/** + * SPI interface for providing custom LeaderNodeSelector implementations. + * Third-party libraries can implement this interface and register their implementations + * using Java's ServiceLoader mechanism. + */ +public interface LeaderNodeSelectorProvider { + + /** + * Creates a new LeaderNodeSelector instance based on the provided configuration. + * + * @return A new LeaderNodeSelector instance + * @throws Exception If the selector cannot be created + */ + LeaderNodeSelector createSelector() throws Exception; +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/RuntimeLeaderRegistry.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/RuntimeLeaderRegistry.java new file mode 100644 index 0000000000..0aa030d46c --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/RuntimeLeaderRegistry.java @@ -0,0 +1,46 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.runtime; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.function.BooleanSupplier; + +/** + * Stores runtime-provided suppliers that answer whether the current process + * should act as the leader. + */ +public final class RuntimeLeaderRegistry { + private static final Logger LOGGER = LoggerFactory.getLogger(RuntimeLeaderRegistry.class); + private static BooleanSupplier supplier = () -> false; + + private RuntimeLeaderRegistry() { + } + + public static void register(BooleanSupplier supplier) { + RuntimeLeaderRegistry.supplier = supplier; + LOGGER.info("Registered runtime leader supplier for log metrics."); + } + + public static BooleanSupplier supplier() { + return supplier; + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/RuntimeLeaderSelectorProvider.java b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/RuntimeLeaderSelectorProvider.java new file mode 100644 index 0000000000..fb28b27b4b --- /dev/null +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/automq/runtime/RuntimeLeaderSelectorProvider.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.connect.automq.runtime; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BooleanSupplier; + +public class RuntimeLeaderSelectorProvider implements LeaderNodeSelectorProvider { + private static final Logger LOGGER = LoggerFactory.getLogger(RuntimeLeaderSelectorProvider.class); + + @Override + public LeaderNodeSelector createSelector() { + final AtomicBoolean missingLogged = new AtomicBoolean(false); + final AtomicBoolean leaderLogged = new AtomicBoolean(false); + + return () -> { + BooleanSupplier current = org.apache.kafka.connect.automq.runtime.RuntimeLeaderRegistry.supplier(); + if (current == null) { + if (missingLogged.compareAndSet(false, true)) { + LOGGER.warn("leader supplier for key not yet available; treating node as follower until registration happens."); + } + if (leaderLogged.getAndSet(false)) { + LOGGER.info("Node stepped down from leadership because supplier is unavailable."); + } + return false; + } + + if (missingLogged.get()) { + missingLogged.set(false); + LOGGER.info("leader supplier is now available."); + } + + try { + boolean leader = current.getAsBoolean(); + if (leader) { + if (!leaderLogged.getAndSet(true)) { + LOGGER.info("Node became leader"); + } + } else { + if (leaderLogged.getAndSet(false)) { + LOGGER.info("Node stepped down from leadership"); + } + } + return leader; + } catch (RuntimeException e) { + if (leaderLogged.getAndSet(false)) { + LOGGER.info("Node stepped down from leadership due to supplier exception."); + } + LOGGER.warn("leader supplier threw exception. Treating as follower.", e); + return false; + } + }; + } +} diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java index 5cfa300baf..dc3419c636 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/AbstractConnectCli.java @@ -19,6 +19,9 @@ import org.apache.kafka.common.utils.Exit; import org.apache.kafka.common.utils.Time; import org.apache.kafka.common.utils.Utils; +import org.apache.kafka.connect.automq.az.AzMetadataProviderHolder; +import org.apache.kafka.connect.automq.log.ConnectLogUploader; +import org.apache.kafka.connect.automq.metrics.OpenTelemetryMetricsReporter; import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy; import org.apache.kafka.connect.runtime.Connect; import org.apache.kafka.connect.runtime.Herder; @@ -36,6 +39,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.Map; +import java.util.Properties; /** * Common initialization logic for Kafka Connect, intended for use by command line utilities @@ -45,7 +49,9 @@ */ public abstract class AbstractConnectCli { - private static final Logger log = LoggerFactory.getLogger(AbstractConnectCli.class); + private static Logger getLogger() { + return LoggerFactory.getLogger(AbstractConnectCli.class); + } private final String[] args; private final Time time = Time.SYSTEM; @@ -83,7 +89,6 @@ protected abstract H createHerder(T config, String workerId, Plugins plugins, */ public void run() { if (args.length < 1 || Arrays.asList(args).contains("--help")) { - log.info("Usage: {}", usage()); Exit.exit(1); } @@ -92,6 +97,17 @@ public void run() { Map workerProps = !workerPropsFile.isEmpty() ? Utils.propsToStringMap(Utils.loadProps(workerPropsFile)) : Collections.emptyMap(); String[] extraArgs = Arrays.copyOfRange(args, 1, args.length); + + // AutoMQ inject start + // Initialize S3 log uploader and OpenTelemetry with worker properties + ConnectLogUploader.initialize(workerProps); + AzMetadataProviderHolder.initialize(workerProps); + + Properties telemetryProps = new Properties(); + telemetryProps.putAll(workerProps); + OpenTelemetryMetricsReporter.initializeTelemetry(telemetryProps); + // AutoMQ inject end + Connect connect = startConnect(workerProps); processExtraArgs(connect, extraArgs); @@ -99,7 +115,7 @@ public void run() { connect.awaitStop(); } catch (Throwable t) { - log.error("Stopping due to error", t); + getLogger().error("Stopping due to error", t); Exit.exit(2); } } @@ -111,17 +127,17 @@ public void run() { * @return a started instance of {@link Connect} */ public Connect startConnect(Map workerProps) { - log.info("Kafka Connect worker initializing ..."); + getLogger().info("Kafka Connect worker initializing ..."); long initStart = time.hiResClockMs(); WorkerInfo initInfo = new WorkerInfo(); initInfo.logAll(); - log.info("Scanning for plugin classes. This might take a moment ..."); + getLogger().info("Scanning for plugin classes. This might take a moment ..."); Plugins plugins = new Plugins(workerProps); plugins.compareAndSwapWithDelegatingLoader(); T config = createConfig(workerProps); - log.debug("Kafka cluster ID: {}", config.kafkaClusterId()); + getLogger().debug("Kafka cluster ID: {}", config.kafkaClusterId()); RestClient restClient = new RestClient(config); @@ -138,11 +154,11 @@ public Connect startConnect(Map workerProps) { H herder = createHerder(config, workerId, plugins, connectorClientConfigOverridePolicy, restServer, restClient); final Connect connect = new Connect<>(herder, restServer); - log.info("Kafka Connect worker initialization took {}ms", time.hiResClockMs() - initStart); + getLogger().info("Kafka Connect worker initialization took {}ms", time.hiResClockMs() - initStart); try { connect.start(); } catch (Exception e) { - log.error("Failed to start Connect", e); + getLogger().error("Failed to start Connect", e); connect.stop(); Exit.exit(3); } diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java index 8763dd908a..84f07fce77 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/cli/ConnectDistributed.java @@ -17,6 +17,7 @@ package org.apache.kafka.connect.cli; import org.apache.kafka.common.utils.Time; +import org.apache.kafka.connect.automq.runtime.RuntimeLeaderRegistry; import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy; import org.apache.kafka.connect.json.JsonConverter; import org.apache.kafka.connect.json.JsonConverterConfig; @@ -39,6 +40,7 @@ import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.function.BooleanSupplier; import static org.apache.kafka.clients.CommonClientConfigs.CLIENT_ID_CONFIG; @@ -96,10 +98,16 @@ protected DistributedHerder createHerder(DistributedConfig config, String worker // Pass the shared admin to the distributed herder as an additional AutoCloseable object that should be closed when the // herder is stopped. This is easier than having to track and own the lifecycle ourselves. - return new DistributedHerder(config, Time.SYSTEM, worker, + DistributedHerder herder = new DistributedHerder(config, Time.SYSTEM, worker, kafkaClusterId, statusBackingStore, configBackingStore, restServer.advertisedUrl().toString(), restClient, connectorClientConfigOverridePolicy, Collections.emptyList(), sharedAdmin); + // AutoMQ for Kafka connect inject start + BooleanSupplier leaderSupplier = herder::isLeaderInstance; + RuntimeLeaderRegistry.register(leaderSupplier); + // AutoMQ for Kafka connect inject end + + return herder; } @Override diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java index d5de59f6a2..88b165eeee 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Connect.java @@ -21,6 +21,8 @@ import org.apache.kafka.connect.runtime.rest.ConnectRestServer; import org.apache.kafka.connect.runtime.rest.RestServer; +import com.automq.log.S3RollingFileAppender; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -115,6 +117,9 @@ public void run() { try { startLatch.await(); Connect.this.stop(); + // AutoMQ inject start + S3RollingFileAppender.shutdown(); + // AutoMQ inject end } catch (InterruptedException e) { log.error("Interrupted in shutdown hook while waiting for Kafka Connect startup to finish"); } diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java index 0a44028a30..08bff7579d 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/Worker.java @@ -48,6 +48,7 @@ import org.apache.kafka.common.utils.Time; import org.apache.kafka.common.utils.Timer; import org.apache.kafka.common.utils.Utils; +import org.apache.kafka.connect.automq.az.AzAwareClientConfigurator; import org.apache.kafka.connect.connector.ConnectRecord; import org.apache.kafka.connect.connector.Connector; import org.apache.kafka.connect.connector.Task; @@ -841,6 +842,10 @@ static Map baseProducerConfigs(String connName, connectorClientConfigOverridePolicy); producerProps.putAll(producerOverrides); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyProducerAz(producerProps, defaultClientId); + // AutoMQ for Kafka inject end + return producerProps; } @@ -909,6 +914,10 @@ static Map baseConsumerConfigs(String connName, connectorClientConfigOverridePolicy); consumerProps.putAll(consumerOverrides); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyConsumerAz(consumerProps, defaultClientId); + // AutoMQ for Kafka inject end + return consumerProps; } @@ -938,6 +947,10 @@ static Map adminConfigs(String connName, // Admin client-specific overrides in the worker config adminProps.putAll(config.originalsWithPrefix("admin.")); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyAdminAz(adminProps, defaultClientId); + // AutoMQ for Kafka inject end + // Connector-specified overrides Map adminOverrides = connectorClientConfigOverrides(connName, connConfig, connectorClass, ConnectorConfig.CONNECTOR_CLIENT_ADMIN_OVERRIDES_PREFIX, diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java index 174adea8cd..2ec2399c5a 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java @@ -1735,6 +1735,12 @@ public void setClusterLoggerLevel(String namespace, String level) { configBackingStore.putLoggerLevel(namespace, level); } + // AutoMQ inject start + public boolean isLeaderInstance() { + return isLeader(); + } + // AutoMQ inject end + // Should only be called from work thread, so synchronization should not be needed protected boolean isLeader() { return assignment != null && member.memberId().equals(assignment.leader()); diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java index 16ccf22f22..a1bd3d4c67 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaConfigBackingStore.java @@ -35,6 +35,7 @@ import org.apache.kafka.common.utils.Time; import org.apache.kafka.common.utils.Timer; import org.apache.kafka.common.utils.Utils; +import org.apache.kafka.connect.automq.az.AzAwareClientConfigurator; import org.apache.kafka.connect.data.Schema; import org.apache.kafka.connect.data.SchemaAndValue; import org.apache.kafka.connect.data.SchemaBuilder; @@ -440,6 +441,9 @@ Map fencableProducerProps(DistributedConfig workerConfig) { Map result = new HashMap<>(baseProducerProps(workerConfig)); result.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId + "-leader"); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyProducerAz(result, "config-log-leader"); + // AutoMQ for Kafka inject end // Always require producer acks to all to ensure durable writes result.put(ProducerConfig.ACKS_CONFIG, "all"); // We can set this to 5 instead of 1 without risking reordering because we are using an idempotent producer @@ -773,11 +777,17 @@ KafkaBasedLog setupAndCreateKafkaBasedLog(String topic, final Wo Map producerProps = new HashMap<>(baseProducerProps); producerProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyProducerAz(producerProps, "config-log"); + // AutoMQ for Kafka inject end Map consumerProps = new HashMap<>(originals); consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName()); consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName()); consumerProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyConsumerAz(consumerProps, "config-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId); if (config.exactlyOnceSourceEnabled()) { ConnectUtils.ensureProperty( @@ -790,6 +800,9 @@ KafkaBasedLog setupAndCreateKafkaBasedLog(String topic, final Wo Map adminProps = new HashMap<>(originals); ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId); adminProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyAdminAz(adminProps, "config-log"); + // AutoMQ for Kafka inject end Map topicSettings = config instanceof DistributedConfig ? ((DistributedConfig) config).configStorageTopicSettings() diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java index 96da411a27..73d37db7c3 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaOffsetBackingStore.java @@ -30,6 +30,7 @@ import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.apache.kafka.common.serialization.ByteArraySerializer; import org.apache.kafka.common.utils.Time; +import org.apache.kafka.connect.automq.az.AzAwareClientConfigurator; import org.apache.kafka.connect.errors.ConnectException; import org.apache.kafka.connect.runtime.WorkerConfig; import org.apache.kafka.connect.runtime.distributed.DistributedConfig; @@ -192,12 +193,18 @@ public void configure(final WorkerConfig config) { // gets approved and scheduled for release. producerProps.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "false"); producerProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyProducerAz(producerProps, "offset-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(producerProps, config, clusterId); Map consumerProps = new HashMap<>(originals); consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName()); consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName()); consumerProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyConsumerAz(consumerProps, "offset-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId); if (config.exactlyOnceSourceEnabled()) { ConnectUtils.ensureProperty( @@ -209,6 +216,9 @@ public void configure(final WorkerConfig config) { Map adminProps = new HashMap<>(originals); adminProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyAdminAz(adminProps, "offset-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId); NewTopic topicDescription = newTopicDescription(topic, config); diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java index 0a9e383700..2bd139f01d 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java @@ -30,6 +30,7 @@ import org.apache.kafka.common.serialization.StringSerializer; import org.apache.kafka.common.utils.ThreadUtils; import org.apache.kafka.common.utils.Time; +import org.apache.kafka.connect.automq.az.AzAwareClientConfigurator; import org.apache.kafka.connect.data.Schema; import org.apache.kafka.connect.data.SchemaAndValue; import org.apache.kafka.connect.data.SchemaBuilder; @@ -183,16 +184,25 @@ public void configure(final WorkerConfig config) { // gets approved and scheduled for release. producerProps.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "false"); // disable idempotence since retries is force to 0 producerProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyProducerAz(producerProps, "status-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(producerProps, config, clusterId); Map consumerProps = new HashMap<>(originals); consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName()); consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName()); consumerProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyConsumerAz(consumerProps, "status-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId); Map adminProps = new HashMap<>(originals); adminProps.put(CommonClientConfigs.CLIENT_ID_CONFIG, clientId); + // AutoMQ for Kafka inject start + AzAwareClientConfigurator.maybeApplyAdminAz(adminProps, "status-log"); + // AutoMQ for Kafka inject end ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId); Map topicSettings = config instanceof DistributedConfig diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/automq/AzAwareClientConfiguratorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/automq/AzAwareClientConfiguratorTest.java new file mode 100644 index 0000000000..07a4f940b9 --- /dev/null +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/automq/AzAwareClientConfiguratorTest.java @@ -0,0 +1,115 @@ +package org.apache.kafka.connect.automq; + +import org.apache.kafka.clients.admin.AdminClientConfig; +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.connect.automq.az.AzAwareClientConfigurator; +import org.apache.kafka.connect.automq.az.AzMetadataProvider; +import org.apache.kafka.connect.automq.az.AzMetadataProviderHolder; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; + +class AzAwareClientConfiguratorTest { + + @AfterEach + void resetProvider() { + AzMetadataProviderHolder.setProviderForTest(null); + } + + @Test + void shouldDecorateProducerClientId() { + AzMetadataProviderHolder.setProviderForTest(new FixedAzProvider("us-east-1a")); + Map props = new HashMap<>(); + props.put(ProducerConfig.CLIENT_ID_CONFIG, "producer-1"); + + AzAwareClientConfigurator.maybeApplyProducerAz(props, "producer-1"); + + assertEquals("automq_type=producer&automq_role=producer-1&automq_az=us-east-1a&producer-1", + props.get(ProducerConfig.CLIENT_ID_CONFIG)); + } + + @Test + void shouldPreserveCustomClientIdInAzConfig() { + AzMetadataProviderHolder.setProviderForTest(new FixedAzProvider("us-east-1a")); + Map props = new HashMap<>(); + props.put(ProducerConfig.CLIENT_ID_CONFIG, "custom-id"); + + AzAwareClientConfigurator.maybeApplyProducerAz(props, "producer-1"); + + assertEquals("automq_type=producer&automq_role=producer-1&automq_az=us-east-1a&custom-id", + props.get(ProducerConfig.CLIENT_ID_CONFIG)); + } + + @Test + void shouldAssignRackForConsumers() { + AzMetadataProviderHolder.setProviderForTest(new FixedAzProvider("us-west-2c")); + Map props = new HashMap<>(); + props.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer-1"); + + AzAwareClientConfigurator.maybeApplyConsumerAz(props, "consumer-1"); + + assertEquals("us-west-2c", props.get(ConsumerConfig.CLIENT_RACK_CONFIG)); + } + + @Test + void shouldDecorateAdminClientId() { + AzMetadataProviderHolder.setProviderForTest(new FixedAzProvider("eu-west-1b")); + Map props = new HashMap<>(); + props.put(AdminClientConfig.CLIENT_ID_CONFIG, "admin-1"); + + AzAwareClientConfigurator.maybeApplyAdminAz(props, "admin-1"); + + assertEquals("automq_type=admin&automq_role=admin-1&automq_az=eu-west-1b&admin-1", + props.get(AdminClientConfig.CLIENT_ID_CONFIG)); + } + + @Test + void shouldLeaveClientIdWhenAzUnavailable() { + AzMetadataProviderHolder.setProviderForTest(new AzMetadataProvider() { + @Override + public Optional availabilityZoneId() { + return Optional.empty(); + } + }); + Map props = new HashMap<>(); + props.put(ProducerConfig.CLIENT_ID_CONFIG, "producer-1"); + + AzAwareClientConfigurator.maybeApplyProducerAz(props, "producer-1"); + + assertEquals("producer-1", props.get(ProducerConfig.CLIENT_ID_CONFIG)); + assertFalse(props.containsKey(ConsumerConfig.CLIENT_RACK_CONFIG)); + } + + @Test + void shouldEncodeSpecialCharactersInClientId() { + AzMetadataProviderHolder.setProviderForTest(new FixedAzProvider("us-east-1a")); + Map props = new HashMap<>(); + props.put(ProducerConfig.CLIENT_ID_CONFIG, "client-with-spaces & symbols"); + + AzAwareClientConfigurator.maybeApplyProducerAz(props, "test-role"); + + assertEquals("automq_type=producer&automq_role=test-role&automq_az=us-east-1a&client-with-spaces & symbols", + props.get(ProducerConfig.CLIENT_ID_CONFIG)); + } + + private static final class FixedAzProvider implements AzMetadataProvider { + private final String az; + + private FixedAzProvider(String az) { + this.az = az; + } + + @Override + public Optional availabilityZoneId() { + return Optional.ofNullable(az); + } + } +} diff --git a/container/automq_chart_feature.patch b/container/automq_chart_feature.patch new file mode 100644 index 0000000000..1135274efd --- /dev/null +++ b/container/automq_chart_feature.patch @@ -0,0 +1,133 @@ +diff --git a/container/bitnami/Dockerfile b/container/bitnami/Dockerfile +index 717a36c21f..ea5eb74efb 100644 +--- a/container/bitnami/Dockerfile ++++ b/container/bitnami/Dockerfile +@@ -1,21 +1,25 @@ + # Copyright Broadcom, Inc. All Rights Reserved. + # SPDX-License-Identifier: APACHE-2.0 + ++FROM docker.io/bitnami/minideb:bookworm as extractor ++ ++COPY automq-*.tgz /tmp/ ++RUN mkdir -p /output && \ ++ tar -zxf /tmp/automq-*.tgz -C /output --strip-components=1 ++ + FROM docker.io/bitnami/minideb:bookworm + + ARG DOWNLOADS_URL="downloads.bitnami.com/files/stacksmith" + ARG JAVA_EXTRA_SECURITY_DIR="/bitnami/java/extra-security" + ARG TARGETARCH + +-LABEL com.vmware.cp.artifact.flavor="sha256:c50c90cfd9d12b445b011e6ad529f1ad3daea45c26d20b00732fae3cd71f6a83" \ +- org.opencontainers.image.base.name="docker.io/bitnami/minideb:bookworm" \ +- org.opencontainers.image.created="2025-03-15T19:51:22Z" \ +- org.opencontainers.image.description="Application packaged by Broadcom, Inc." \ +- org.opencontainers.image.documentation="https://github.com/bitnami/containers/tree/main/bitnami/kafka/README.md" \ ++LABEL org.opencontainers.image.base.name="docker.io/bitnami/minideb:bookworm" \ ++ org.opencontainers.image.created="2025-04-27T21:51:40Z" \ ++ org.opencontainers.image.description="AutoMQ packaged by AutoMQ, Inc." \ + org.opencontainers.image.ref.name="3.9.0-debian-12-r13" \ +- org.opencontainers.image.source="https://github.com/bitnami/containers/tree/main/bitnami/kafka" \ +- org.opencontainers.image.title="kafka" \ +- org.opencontainers.image.vendor="Broadcom, Inc." \ ++ org.opencontainers.image.source="https://github.com/AutoMQ/automq/tree/main/container" \ ++ org.opencontainers.image.title="automq" \ ++ org.opencontainers.image.vendor="AutoMQ, Inc." \ + org.opencontainers.image.version="3.9.0" + + ENV HOME="/" \ +@@ -26,12 +30,11 @@ ENV HOME="/" \ + COPY prebuildfs / + SHELL ["/bin/bash", "-o", "errexit", "-o", "nounset", "-o", "pipefail", "-c"] + # Install required system packages and dependencies +-RUN install_packages ca-certificates curl procps zlib1g +-RUN mkdir -p /tmp/bitnami/pkg/cache/ ; cd /tmp/bitnami/pkg/cache/ ; \ ++RUN apt-get update && apt-get install -y ca-certificates curl procps zlib1g libjemalloc-dev && \ ++ apt-get clean && rm -rf /var/lib/apt/lists /var/cache/apt/archives ++RUN mkdir -p /tmp/bitnami/pkg/cache/ ; cd /tmp/bitnami/pkg/cache/ || exit 1 ; \ + COMPONENTS=( \ +- "wait-for-port-1.0.8-14-linux-${OS_ARCH}-debian-12" \ + "jre-17.0.14-10-1-linux-${OS_ARCH}-debian-12" \ +- "kafka-3.9.0-2-linux-${OS_ARCH}-debian-12" \ + ) ; \ + for COMPONENT in "${COMPONENTS[@]}"; do \ + if [ ! -f "${COMPONENT}.tar.gz" ]; then \ +@@ -39,17 +42,18 @@ RUN mkdir -p /tmp/bitnami/pkg/cache/ ; cd /tmp/bitnami/pkg/cache/ ; \ + curl -SsLf "https://${DOWNLOADS_URL}/${COMPONENT}.tar.gz.sha256" -O ; \ + fi ; \ + sha256sum -c "${COMPONENT}.tar.gz.sha256" ; \ +- tar -zxf "${COMPONENT}.tar.gz" -C /opt/bitnami --strip-components=2 --no-same-owner --wildcards '*/files' ; \ ++ tar -zxf "${COMPONENT}.tar.gz" -C /opt/bitnami --strip-components=2 --no-same-owner ; \ + rm -rf "${COMPONENT}".tar.gz{,.sha256} ; \ + done +-RUN apt-get update && apt-get upgrade -y && \ +- apt-get clean && rm -rf /var/lib/apt/lists /var/cache/apt/archives ++COPY --from=extractor --chown=1001:0 /output /opt/bitnami/kafka + RUN chmod g+rwX /opt/bitnami + RUN find / -perm /6000 -type f -exec chmod a-s {} \; || true + RUN ln -s /opt/bitnami/scripts/kafka/entrypoint.sh /entrypoint.sh + RUN ln -s /opt/bitnami/scripts/kafka/run.sh /run.sh + + COPY rootfs / ++RUN find /opt/bitnami/scripts -type f -exec chmod g+rwX {} \; && \ ++ find /opt/bitnami/scripts -type f -exec chmod +x {} \; + RUN /opt/bitnami/scripts/java/postunpack.sh + RUN /opt/bitnami/scripts/kafka/postunpack.sh + ENV APP_VERSION="3.9.0" \ +@@ -59,6 +63,12 @@ ENV APP_VERSION="3.9.0" \ + + EXPOSE 9092 + ++RUN arch=$(uname -m) \ ++ && target_file="/usr/lib/${arch}-linux-gnu/libjemalloc.so" \ ++ && { test -f "$target_file" || { echo "Error: $target_file not found"; exit 1; }; } \ ++ && ln -sv "$target_file" /usr/lib/libjemalloc.so ++ENV LD_PRELOAD="/usr/lib/libjemalloc.so" ++ + USER 1001 + ENTRYPOINT [ "/opt/bitnami/scripts/kafka/entrypoint.sh" ] + CMD [ "/opt/bitnami/scripts/kafka/run.sh" ] +diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh +index 00d053b521..09e3d3084d 100644 +--- a/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh ++++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh +@@ -42,12 +42,13 @@ print_welcome_page() { + # None + ######################### + print_image_welcome_page() { +- local github_url="https://github.com/bitnami/containers" ++ local docs_url="https://www.automq.com/docs/automq/deployment/deploy-multi-nodes-cluster-on-kubernetes" + + info "" +- info "${BOLD}Welcome to the Bitnami ${BITNAMI_APP_NAME} container${RESET}" +- info "Subscribe to project updates by watching ${BOLD}${github_url}${RESET}" +- info "Did you know there are enterprise versions of the Bitnami catalog? For enhanced secure software supply chain features, unlimited pulls from Docker, LTS support, or application customization, see Bitnami Premium or Tanzu Application Catalog. See https://www.arrow.com/globalecs/na/vendors/bitnami/ for more information." ++ info "${BOLD}Welcome to the AutoMQ for Apache Kafka on Bitnami Container${RESET}" ++ info "${BOLD}This image is compatible with Bitnami's container standards.${RESET}" ++ info "Refer to the documentation for complete configuration and Kubernetes deployment guidelines:" ++ info "${BOLD}${docs_url}${RESET}" + info "" + } + +diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh +index 7255563236..673c84e721 100644 +--- a/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh ++++ b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh +@@ -35,12 +35,12 @@ mv "${KAFKA_CONF_DIR}/server.properties" "${KAFKA_CONF_DIR}/server.properties.or + + # Disable logging to stdout and garbage collection + # Source: https://logging.apache.org/log4j/log4j-2.4/manual/appenders.html +-replace_in_file "${KAFKA_BASE_DIR}/bin/kafka-server-start.sh" " [-]loggc" " " +-replace_in_file "${KAFKA_CONF_DIR}/log4j.properties" "DailyRollingFileAppender" "ConsoleAppender" ++#replace_in_file "${KAFKA_BASE_DIR}/bin/kafka-server-start.sh" " [-]loggc" " " ++#replace_in_file "${KAFKA_CONF_DIR}/log4j.properties" "DailyRollingFileAppender" "ConsoleAppender" + + # Disable the default console logger in favour of KafkaAppender (which provides the exact output) +-echo "log4j.appender.stdout.Threshold=OFF" >>"${KAFKA_CONF_DIR}/log4j.properties" ++#echo "log4j.appender.stdout.Threshold=OFF" >>"${KAFKA_CONF_DIR}/log4j.properties" + + # Remove invalid parameters for ConsoleAppender +-remove_in_file "${KAFKA_CONF_DIR}/log4j.properties" "DatePattern" +-remove_in_file "${KAFKA_CONF_DIR}/log4j.properties" "Appender.File" ++#remove_in_file "${KAFKA_CONF_DIR}/log4j.properties" "DatePattern" ++#remove_in_file "${KAFKA_CONF_DIR}/log4j.properties" "Appender.File" diff --git a/container/bitnami/Dockerfile b/container/bitnami/Dockerfile new file mode 100644 index 0000000000..ea5eb74efb --- /dev/null +++ b/container/bitnami/Dockerfile @@ -0,0 +1,74 @@ +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +FROM docker.io/bitnami/minideb:bookworm as extractor + +COPY automq-*.tgz /tmp/ +RUN mkdir -p /output && \ + tar -zxf /tmp/automq-*.tgz -C /output --strip-components=1 + +FROM docker.io/bitnami/minideb:bookworm + +ARG DOWNLOADS_URL="downloads.bitnami.com/files/stacksmith" +ARG JAVA_EXTRA_SECURITY_DIR="/bitnami/java/extra-security" +ARG TARGETARCH + +LABEL org.opencontainers.image.base.name="docker.io/bitnami/minideb:bookworm" \ + org.opencontainers.image.created="2025-04-27T21:51:40Z" \ + org.opencontainers.image.description="AutoMQ packaged by AutoMQ, Inc." \ + org.opencontainers.image.ref.name="3.9.0-debian-12-r13" \ + org.opencontainers.image.source="https://github.com/AutoMQ/automq/tree/main/container" \ + org.opencontainers.image.title="automq" \ + org.opencontainers.image.vendor="AutoMQ, Inc." \ + org.opencontainers.image.version="3.9.0" + +ENV HOME="/" \ + OS_ARCH="${TARGETARCH:-amd64}" \ + OS_FLAVOUR="debian-12" \ + OS_NAME="linux" + +COPY prebuildfs / +SHELL ["/bin/bash", "-o", "errexit", "-o", "nounset", "-o", "pipefail", "-c"] +# Install required system packages and dependencies +RUN apt-get update && apt-get install -y ca-certificates curl procps zlib1g libjemalloc-dev && \ + apt-get clean && rm -rf /var/lib/apt/lists /var/cache/apt/archives +RUN mkdir -p /tmp/bitnami/pkg/cache/ ; cd /tmp/bitnami/pkg/cache/ || exit 1 ; \ + COMPONENTS=( \ + "jre-17.0.14-10-1-linux-${OS_ARCH}-debian-12" \ + ) ; \ + for COMPONENT in "${COMPONENTS[@]}"; do \ + if [ ! -f "${COMPONENT}.tar.gz" ]; then \ + curl -SsLf "https://${DOWNLOADS_URL}/${COMPONENT}.tar.gz" -O ; \ + curl -SsLf "https://${DOWNLOADS_URL}/${COMPONENT}.tar.gz.sha256" -O ; \ + fi ; \ + sha256sum -c "${COMPONENT}.tar.gz.sha256" ; \ + tar -zxf "${COMPONENT}.tar.gz" -C /opt/bitnami --strip-components=2 --no-same-owner ; \ + rm -rf "${COMPONENT}".tar.gz{,.sha256} ; \ + done +COPY --from=extractor --chown=1001:0 /output /opt/bitnami/kafka +RUN chmod g+rwX /opt/bitnami +RUN find / -perm /6000 -type f -exec chmod a-s {} \; || true +RUN ln -s /opt/bitnami/scripts/kafka/entrypoint.sh /entrypoint.sh +RUN ln -s /opt/bitnami/scripts/kafka/run.sh /run.sh + +COPY rootfs / +RUN find /opt/bitnami/scripts -type f -exec chmod g+rwX {} \; && \ + find /opt/bitnami/scripts -type f -exec chmod +x {} \; +RUN /opt/bitnami/scripts/java/postunpack.sh +RUN /opt/bitnami/scripts/kafka/postunpack.sh +ENV APP_VERSION="3.9.0" \ + BITNAMI_APP_NAME="kafka" \ + JAVA_HOME="/opt/bitnami/java" \ + PATH="/opt/bitnami/common/bin:/opt/bitnami/java/bin:/opt/bitnami/kafka/bin:$PATH" + +EXPOSE 9092 + +RUN arch=$(uname -m) \ + && target_file="/usr/lib/${arch}-linux-gnu/libjemalloc.so" \ + && { test -f "$target_file" || { echo "Error: $target_file not found"; exit 1; }; } \ + && ln -sv "$target_file" /usr/lib/libjemalloc.so +ENV LD_PRELOAD="/usr/lib/libjemalloc.so" + +USER 1001 +ENTRYPOINT [ "/opt/bitnami/scripts/kafka/entrypoint.sh" ] +CMD [ "/opt/bitnami/scripts/kafka/run.sh" ] diff --git a/container/bitnami/docker-compose.yml b/container/bitnami/docker-compose.yml new file mode 100644 index 0000000000..3cb537c151 --- /dev/null +++ b/container/bitnami/docker-compose.yml @@ -0,0 +1,24 @@ +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +services: + kafka: + image: docker.io/bitnami/kafka:3.9 + ports: + - "9092:9092" + volumes: + - "kafka_data:/bitnami" + environment: + # KRaft settings + - KAFKA_CFG_NODE_ID=0 + - KAFKA_CFG_PROCESS_ROLES=controller,broker + - KAFKA_CFG_CONTROLLER_QUORUM_VOTERS=0@kafka:9093 + # Listeners + - KAFKA_CFG_LISTENERS=PLAINTEXT://:9092,CONTROLLER://:9093 + - KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://:9092 + - KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT + - KAFKA_CFG_CONTROLLER_LISTENER_NAMES=CONTROLLER + - KAFKA_CFG_INTER_BROKER_LISTENER_NAME=PLAINTEXT +volumes: + kafka_data: + driver: local diff --git a/container/bitnami/prebuildfs/opt/bitnami/.bitnami_components.json b/container/bitnami/prebuildfs/opt/bitnami/.bitnami_components.json new file mode 100644 index 0000000000..92767c41dd --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/.bitnami_components.json @@ -0,0 +1,20 @@ +{ + "jre": { + "arch": "amd64", + "distro": "debian-12", + "type": "NAMI", + "version": "17.0.14-10-1" + }, + "kafka": { + "arch": "amd64", + "distro": "debian-12", + "type": "NAMI", + "version": "3.9.0-2" + }, + "wait-for-port": { + "arch": "amd64", + "distro": "debian-12", + "type": "NAMI", + "version": "1.0.8-14" + } +} \ No newline at end of file diff --git a/container/bitnami/prebuildfs/opt/bitnami/licenses/licenses.txt b/container/bitnami/prebuildfs/opt/bitnami/licenses/licenses.txt new file mode 100644 index 0000000000..76956b38e8 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/licenses/licenses.txt @@ -0,0 +1,2 @@ +Bitnami containers ship with software bundles. You can find the licenses under: +/opt/bitnami/[name-of-bundle]/licenses/[bundle-version].txt diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh new file mode 100644 index 0000000000..09e3d3084d --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libbitnami.sh @@ -0,0 +1,54 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Bitnami custom library + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/liblog.sh + +# Constants +BOLD='\033[1m' + +# Functions + +######################## +# Print the welcome page +# Globals: +# DISABLE_WELCOME_MESSAGE +# BITNAMI_APP_NAME +# Arguments: +# None +# Returns: +# None +######################### +print_welcome_page() { + if [[ -z "${DISABLE_WELCOME_MESSAGE:-}" ]]; then + if [[ -n "$BITNAMI_APP_NAME" ]]; then + print_image_welcome_page + fi + fi +} + +######################## +# Print the welcome page for a Bitnami Docker image +# Globals: +# BITNAMI_APP_NAME +# Arguments: +# None +# Returns: +# None +######################### +print_image_welcome_page() { + local docs_url="https://www.automq.com/docs/automq/deployment/deploy-multi-nodes-cluster-on-kubernetes" + + info "" + info "${BOLD}Welcome to the AutoMQ for Apache Kafka on Bitnami Container${RESET}" + info "${BOLD}This image is compatible with Bitnami's container standards.${RESET}" + info "Refer to the documentation for complete configuration and Kubernetes deployment guidelines:" + info "${BOLD}${docs_url}${RESET}" + info "" +} + diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libfile.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libfile.sh new file mode 100644 index 0000000000..1c69e0e48a --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libfile.sh @@ -0,0 +1,141 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for managing files + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/libos.sh + +# Functions + +######################## +# Replace a regex-matching string in a file +# Arguments: +# $1 - filename +# $2 - match regex +# $3 - substitute regex +# $4 - use POSIX regex. Default: true +# Returns: +# None +######################### +replace_in_file() { + local filename="${1:?filename is required}" + local match_regex="${2:?match regex is required}" + local substitute_regex="${3:?substitute regex is required}" + local posix_regex=${4:-true} + + local result + + # We should avoid using 'sed in-place' substitutions + # 1) They are not compatible with files mounted from ConfigMap(s) + # 2) We found incompatibility issues with Debian10 and "in-place" substitutions + local -r del=$'\001' # Use a non-printable character as a 'sed' delimiter to avoid issues + if [[ $posix_regex = true ]]; then + result="$(sed -E "s${del}${match_regex}${del}${substitute_regex}${del}g" "$filename")" + else + result="$(sed "s${del}${match_regex}${del}${substitute_regex}${del}g" "$filename")" + fi + echo "$result" > "$filename" +} + +######################## +# Replace a regex-matching multiline string in a file +# Arguments: +# $1 - filename +# $2 - match regex +# $3 - substitute regex +# Returns: +# None +######################### +replace_in_file_multiline() { + local filename="${1:?filename is required}" + local match_regex="${2:?match regex is required}" + local substitute_regex="${3:?substitute regex is required}" + + local result + local -r del=$'\001' # Use a non-printable character as a 'sed' delimiter to avoid issues + result="$(perl -pe "BEGIN{undef $/;} s${del}${match_regex}${del}${substitute_regex}${del}sg" "$filename")" + echo "$result" > "$filename" +} + +######################## +# Remove a line in a file based on a regex +# Arguments: +# $1 - filename +# $2 - match regex +# $3 - use POSIX regex. Default: true +# Returns: +# None +######################### +remove_in_file() { + local filename="${1:?filename is required}" + local match_regex="${2:?match regex is required}" + local posix_regex=${3:-true} + local result + + # We should avoid using 'sed in-place' substitutions + # 1) They are not compatible with files mounted from ConfigMap(s) + # 2) We found incompatibility issues with Debian10 and "in-place" substitutions + if [[ $posix_regex = true ]]; then + result="$(sed -E "/$match_regex/d" "$filename")" + else + result="$(sed "/$match_regex/d" "$filename")" + fi + echo "$result" > "$filename" +} + +######################## +# Appends text after the last line matching a pattern +# Arguments: +# $1 - file +# $2 - match regex +# $3 - contents to add +# Returns: +# None +######################### +append_file_after_last_match() { + local file="${1:?missing file}" + local match_regex="${2:?missing pattern}" + local value="${3:?missing value}" + + # We read the file in reverse, replace the first match (0,/pattern/s) and then reverse the results again + result="$(tac "$file" | sed -E "0,/($match_regex)/s||${value}\n\1|" | tac)" + echo "$result" > "$file" +} + +######################## +# Wait until certain entry is present in a log file +# Arguments: +# $1 - entry to look for +# $2 - log file +# $3 - max retries. Default: 12 +# $4 - sleep between retries (in seconds). Default: 5 +# Returns: +# Boolean +######################### +wait_for_log_entry() { + local -r entry="${1:-missing entry}" + local -r log_file="${2:-missing log file}" + local -r retries="${3:-12}" + local -r interval_time="${4:-5}" + local attempt=0 + + check_log_file_for_entry() { + if ! grep -qE "$entry" "$log_file"; then + debug "Entry \"${entry}\" still not present in ${log_file} (attempt $((++attempt))/${retries})" + return 1 + fi + } + debug "Checking that ${log_file} log file contains entry \"${entry}\"" + if retry_while check_log_file_for_entry "$retries" "$interval_time"; then + debug "Found entry \"${entry}\" in ${log_file}" + true + else + error "Could not find entry \"${entry}\" in ${log_file} after ${retries} retries" + debug_execute cat "$log_file" + return 1 + fi +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libfs.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libfs.sh new file mode 100644 index 0000000000..1337f6c2dd --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libfs.sh @@ -0,0 +1,193 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for file system actions + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/liblog.sh + +# Functions + +######################## +# Ensure a file/directory is owned (user and group) but the given user +# Arguments: +# $1 - filepath +# $2 - owner +# Returns: +# None +######################### +owned_by() { + local path="${1:?path is missing}" + local owner="${2:?owner is missing}" + local group="${3:-}" + + if [[ -n $group ]]; then + chown "$owner":"$group" "$path" + else + chown "$owner":"$owner" "$path" + fi +} + +######################## +# Ensure a directory exists and, optionally, is owned by the given user +# Arguments: +# $1 - directory +# $2 - owner +# Returns: +# None +######################### +ensure_dir_exists() { + local dir="${1:?directory is missing}" + local owner_user="${2:-}" + local owner_group="${3:-}" + + [ -d "${dir}" ] || mkdir -p "${dir}" + if [[ -n $owner_user ]]; then + owned_by "$dir" "$owner_user" "$owner_group" + fi +} + +######################## +# Checks whether a directory is empty or not +# arguments: +# $1 - directory +# returns: +# boolean +######################### +is_dir_empty() { + local -r path="${1:?missing directory}" + # Calculate real path in order to avoid issues with symlinks + local -r dir="$(realpath "$path")" + if [[ ! -e "$dir" ]] || [[ -z "$(ls -A "$dir")" ]]; then + true + else + false + fi +} + +######################## +# Checks whether a mounted directory is empty or not +# arguments: +# $1 - directory +# returns: +# boolean +######################### +is_mounted_dir_empty() { + local dir="${1:?missing directory}" + + if is_dir_empty "$dir" || find "$dir" -mindepth 1 -maxdepth 1 -not -name ".snapshot" -not -name "lost+found" -exec false {} +; then + true + else + false + fi +} + +######################## +# Checks whether a file can be written to or not +# arguments: +# $1 - file +# returns: +# boolean +######################### +is_file_writable() { + local file="${1:?missing file}" + local dir + dir="$(dirname "$file")" + + if [[ (-f "$file" && -w "$file") || (! -f "$file" && -d "$dir" && -w "$dir") ]]; then + true + else + false + fi +} + +######################## +# Relativize a path +# arguments: +# $1 - path +# $2 - base +# returns: +# None +######################### +relativize() { + local -r path="${1:?missing path}" + local -r base="${2:?missing base}" + pushd "$base" >/dev/null || exit + realpath -q --no-symlinks --relative-base="$base" "$path" | sed -e 's|^/$|.|' -e 's|^/||' + popd >/dev/null || exit +} + +######################## +# Configure permissions and ownership recursively +# Globals: +# None +# Arguments: +# $1 - paths (as a string). +# Flags: +# -f|--file-mode - mode for directories. +# -d|--dir-mode - mode for files. +# -u|--user - user +# -g|--group - group +# Returns: +# None +######################### +configure_permissions_ownership() { + local -r paths="${1:?paths is missing}" + local dir_mode="" + local file_mode="" + local user="" + local group="" + + # Validate arguments + shift 1 + while [ "$#" -gt 0 ]; do + case "$1" in + -f | --file-mode) + shift + file_mode="${1:?missing mode for files}" + ;; + -d | --dir-mode) + shift + dir_mode="${1:?missing mode for directories}" + ;; + -u | --user) + shift + user="${1:?missing user}" + ;; + -g | --group) + shift + group="${1:?missing group}" + ;; + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + + read -r -a filepaths <<<"$paths" + for p in "${filepaths[@]}"; do + if [[ -e "$p" ]]; then + find -L "$p" -printf "" + if [[ -n $dir_mode ]]; then + find -L "$p" -type d ! -perm "$dir_mode" -print0 | xargs -r -0 chmod "$dir_mode" + fi + if [[ -n $file_mode ]]; then + find -L "$p" -type f ! -perm "$file_mode" -print0 | xargs -r -0 chmod "$file_mode" + fi + if [[ -n $user ]] && [[ -n $group ]]; then + find -L "$p" -print0 | xargs -r -0 chown "${user}:${group}" + elif [[ -n $user ]] && [[ -z $group ]]; then + find -L "$p" -print0 | xargs -r -0 chown "${user}" + elif [[ -z $user ]] && [[ -n $group ]]; then + find -L "$p" -print0 | xargs -r -0 chgrp "${group}" + fi + else + stderr_print "$p does not exist" + fi + done +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libhook.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libhook.sh new file mode 100644 index 0000000000..f3a5fe7868 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libhook.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library to use for scripts expected to be used as Kubernetes lifecycle hooks + +# shellcheck disable=SC1091 + +# Load generic libraries +. /opt/bitnami/scripts/liblog.sh +. /opt/bitnami/scripts/libos.sh + +# Override functions that log to stdout/stderr of the current process, so they print to process 1 +for function_to_override in stderr_print debug_execute; do + # Output is sent to output of process 1 and thus end up in the container log + # The hook output in general isn't saved + eval "$(declare -f "$function_to_override") >/proc/1/fd/1 2>/proc/1/fd/2" +done diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/liblog.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/liblog.sh new file mode 100644 index 0000000000..1e18ef9866 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/liblog.sh @@ -0,0 +1,146 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for logging functions + +# Constants +RESET='\033[0m' +RED='\033[38;5;1m' +GREEN='\033[38;5;2m' +YELLOW='\033[38;5;3m' +MAGENTA='\033[38;5;5m' +CYAN='\033[38;5;6m' + +# Functions + +######################## +# Print to STDERR +# Arguments: +# Message to print +# Returns: +# None +######################### +stderr_print() { + # 'is_boolean_yes' is defined in libvalidations.sh, but depends on this file so we cannot source it + local bool="${BITNAMI_QUIET:-false}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if ! [[ "$bool" = 1 || "$bool" =~ ^(yes|true)$ ]]; then + printf "%b\\n" "${*}" >&2 + fi +} + +######################## +# Log message +# Arguments: +# Message to log +# Returns: +# None +######################### +log() { + local color_bool="${BITNAMI_COLOR:-true}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if [[ "$color_bool" = 1 || "$color_bool" =~ ^(yes|true)$ ]]; then + stderr_print "${CYAN}${MODULE:-} ${MAGENTA}$(date "+%T.%2N ")${RESET}${*}" + else + stderr_print "${MODULE:-} $(date "+%T.%2N ")${*}" + fi +} +######################## +# Log an 'info' message +# Arguments: +# Message to log +# Returns: +# None +######################### +info() { + local msg_color="" + local color_bool="${BITNAMI_COLOR:-true}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if [[ "$color_bool" = 1 || "$color_bool" =~ ^(yes|true)$ ]];then + msg_color="$GREEN" + fi + log "${msg_color}INFO ${RESET} ==> ${*}" +} +######################## +# Log message +# Arguments: +# Message to log +# Returns: +# None +######################### +warn() { + local msg_color="" + local color_bool="${BITNAMI_COLOR:-true}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if [[ "$color_bool" = 1 || "$color_bool" =~ ^(yes|true)$ ]];then + msg_color="$YELLOW" + fi + log "${msg_color}WARN ${RESET} ==> ${*}" +} +######################## +# Log an 'error' message +# Arguments: +# Message to log +# Returns: +# None +######################### +error() { + local msg_color="" + local color_bool="${BITNAMI_COLOR:-true}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if [[ "$color_bool" = 1 || "$color_bool" =~ ^(yes|true)$ ]];then + msg_color="$RED" + fi + log "${msg_color}ERROR${RESET} ==> ${*}" +} +######################## +# Log a 'debug' message +# Globals: +# BITNAMI_DEBUG +# Arguments: +# None +# Returns: +# None +######################### +debug() { + local msg_color="" + local color_bool="${BITNAMI_COLOR:-true}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if [[ "$color_bool" = 1 || "$color_bool" =~ ^(yes|true)$ ]] ;then + msg_color="$MAGENTA" + fi + local debug_bool="${BITNAMI_DEBUG:-false}" + if [[ "$debug_bool" = 1 || "$debug_bool" =~ ^(yes|true)$ ]]; then + log "${msg_color}DEBUG${RESET} ==> ${*}" + fi +} + +######################## +# Indent a string +# Arguments: +# $1 - string +# $2 - number of indentation characters (default: 4) +# $3 - indentation character (default: " ") +# Returns: +# None +######################### +indent() { + local string="${1:-}" + local num="${2:?missing num}" + local char="${3:-" "}" + # Build the indentation unit string + local indent_unit="" + for ((i = 0; i < num; i++)); do + indent_unit="${indent_unit}${char}" + done + # shellcheck disable=SC2001 + # Complex regex, see https://github.com/koalaman/shellcheck/wiki/SC2001#exceptions + echo "$string" | sed "s/^/${indent_unit}/" +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libnet.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libnet.sh new file mode 100644 index 0000000000..004e426fba --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libnet.sh @@ -0,0 +1,171 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for network functions + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/liblog.sh +. /opt/bitnami/scripts/libvalidations.sh + +# Functions + +######################## +# Resolve IP address for a host/domain (i.e. DNS lookup) +# Arguments: +# $1 - Hostname to resolve +# $2 - IP address version (v4, v6), leave empty for resolving to any version +# Returns: +# IP +######################### +dns_lookup() { + local host="${1:?host is missing}" + local ip_version="${2:-}" + getent "ahosts${ip_version}" "$host" | awk '/STREAM/ {print $1 }' | head -n 1 +} + +######################### +# Wait for a hostname and return the IP +# Arguments: +# $1 - hostname +# $2 - number of retries +# $3 - seconds to wait between retries +# Returns: +# - IP address that corresponds to the hostname +######################### +wait_for_dns_lookup() { + local hostname="${1:?hostname is missing}" + local retries="${2:-5}" + local seconds="${3:-1}" + check_host() { + if [[ $(dns_lookup "$hostname") == "" ]]; then + false + else + true + fi + } + # Wait for the host to be ready + retry_while "check_host ${hostname}" "$retries" "$seconds" + dns_lookup "$hostname" +} + +######################## +# Get machine's IP +# Arguments: +# None +# Returns: +# Machine IP +######################### +get_machine_ip() { + local -a ip_addresses + local hostname + hostname="$(hostname)" + read -r -a ip_addresses <<< "$(dns_lookup "$hostname" | xargs echo)" + if [[ "${#ip_addresses[@]}" -gt 1 ]]; then + warn "Found more than one IP address associated to hostname ${hostname}: ${ip_addresses[*]}, will use ${ip_addresses[0]}" + elif [[ "${#ip_addresses[@]}" -lt 1 ]]; then + error "Could not find any IP address associated to hostname ${hostname}" + exit 1 + fi + # Check if the first IP address is IPv6 to add brackets + if validate_ipv6 "${ip_addresses[0]}" ; then + echo "[${ip_addresses[0]}]" + else + echo "${ip_addresses[0]}" + fi +} + +######################## +# Check if the provided argument is a resolved hostname +# Arguments: +# $1 - Value to check +# Returns: +# Boolean +######################### +is_hostname_resolved() { + local -r host="${1:?missing value}" + if [[ -n "$(dns_lookup "$host")" ]]; then + true + else + false + fi +} + +######################## +# Parse URL +# Globals: +# None +# Arguments: +# $1 - uri - String +# $2 - component to obtain. Valid options (scheme, authority, userinfo, host, port, path, query or fragment) - String +# Returns: +# String +parse_uri() { + local uri="${1:?uri is missing}" + local component="${2:?component is missing}" + + # Solution based on https://tools.ietf.org/html/rfc3986#appendix-B with + # additional sub-expressions to split authority into userinfo, host and port + # Credits to Patryk Obara (see https://stackoverflow.com/a/45977232/6694969) + local -r URI_REGEX='^(([^:/?#]+):)?(//((([^@/?#]+)@)?([^:/?#]+)(:([0-9]+))?))?(/([^?#]*))?(\?([^#]*))?(#(.*))?' + # || | ||| | | | | | | | | | + # |2 scheme | ||6 userinfo 7 host | 9 port | 11 rpath | 13 query | 15 fragment + # 1 scheme: | |5 userinfo@ 8 :... 10 path 12 ?... 14 #... + # | 4 authority + # 3 //... + local index=0 + case "$component" in + scheme) + index=2 + ;; + authority) + index=4 + ;; + userinfo) + index=6 + ;; + host) + index=7 + ;; + port) + index=9 + ;; + path) + index=10 + ;; + query) + index=13 + ;; + fragment) + index=14 + ;; + *) + stderr_print "unrecognized component $component" + return 1 + ;; + esac + [[ "$uri" =~ $URI_REGEX ]] && echo "${BASH_REMATCH[${index}]}" +} + +######################## +# Wait for a HTTP connection to succeed +# Globals: +# * +# Arguments: +# $1 - URL to wait for +# $2 - Maximum amount of retries (optional) +# $3 - Time between retries (optional) +# Returns: +# true if the HTTP connection succeeded, false otherwise +######################### +wait_for_http_connection() { + local url="${1:?missing url}" + local retries="${2:-}" + local sleep_time="${3:-}" + if ! retry_while "debug_execute curl --silent ${url}" "$retries" "$sleep_time"; then + error "Could not connect to ${url}" + return 1 + fi +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libos.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libos.sh new file mode 100644 index 0000000000..9d908c4857 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libos.sh @@ -0,0 +1,657 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for operating system actions + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/liblog.sh +. /opt/bitnami/scripts/libfs.sh +. /opt/bitnami/scripts/libvalidations.sh + +# Functions + +######################## +# Check if an user exists in the system +# Arguments: +# $1 - user +# Returns: +# Boolean +######################### +user_exists() { + local user="${1:?user is missing}" + id "$user" >/dev/null 2>&1 +} + +######################## +# Check if a group exists in the system +# Arguments: +# $1 - group +# Returns: +# Boolean +######################### +group_exists() { + local group="${1:?group is missing}" + getent group "$group" >/dev/null 2>&1 +} + +######################## +# Create a group in the system if it does not exist already +# Arguments: +# $1 - group +# Flags: +# -i|--gid - the ID for the new group +# -s|--system - Whether to create new user as system user (uid <= 999) +# Returns: +# None +######################### +ensure_group_exists() { + local group="${1:?group is missing}" + local gid="" + local is_system_user=false + + # Validate arguments + shift 1 + while [ "$#" -gt 0 ]; do + case "$1" in + -i | --gid) + shift + gid="${1:?missing gid}" + ;; + -s | --system) + is_system_user=true + ;; + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + + if ! group_exists "$group"; then + local -a args=("$group") + if [[ -n "$gid" ]]; then + if group_exists "$gid"; then + error "The GID $gid is already in use." >&2 + return 1 + fi + args+=("--gid" "$gid") + fi + $is_system_user && args+=("--system") + groupadd "${args[@]}" >/dev/null 2>&1 + fi +} + +######################## +# Create an user in the system if it does not exist already +# Arguments: +# $1 - user +# Flags: +# -i|--uid - the ID for the new user +# -g|--group - the group the new user should belong to +# -a|--append-groups - comma-separated list of supplemental groups to append to the new user +# -h|--home - the home directory for the new user +# -s|--system - whether to create new user as system user (uid <= 999) +# Returns: +# None +######################### +ensure_user_exists() { + local user="${1:?user is missing}" + local uid="" + local group="" + local append_groups="" + local home="" + local is_system_user=false + + # Validate arguments + shift 1 + while [ "$#" -gt 0 ]; do + case "$1" in + -i | --uid) + shift + uid="${1:?missing uid}" + ;; + -g | --group) + shift + group="${1:?missing group}" + ;; + -a | --append-groups) + shift + append_groups="${1:?missing append_groups}" + ;; + -h | --home) + shift + home="${1:?missing home directory}" + ;; + -s | --system) + is_system_user=true + ;; + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + + if ! user_exists "$user"; then + local -a user_args=("-N" "$user") + if [[ -n "$uid" ]]; then + if user_exists "$uid"; then + error "The UID $uid is already in use." + return 1 + fi + user_args+=("--uid" "$uid") + else + $is_system_user && user_args+=("--system") + fi + useradd "${user_args[@]}" >/dev/null 2>&1 + fi + + if [[ -n "$group" ]]; then + local -a group_args=("$group") + $is_system_user && group_args+=("--system") + ensure_group_exists "${group_args[@]}" + usermod -g "$group" "$user" >/dev/null 2>&1 + fi + + if [[ -n "$append_groups" ]]; then + local -a groups + read -ra groups <<<"$(tr ',;' ' ' <<<"$append_groups")" + for group in "${groups[@]}"; do + ensure_group_exists "$group" + usermod -aG "$group" "$user" >/dev/null 2>&1 + done + fi + + if [[ -n "$home" ]]; then + mkdir -p "$home" + usermod -d "$home" "$user" >/dev/null 2>&1 + configure_permissions_ownership "$home" -d "775" -f "664" -u "$user" -g "$group" + fi +} + +######################## +# Check if the script is currently running as root +# Arguments: +# $1 - user +# $2 - group +# Returns: +# Boolean +######################### +am_i_root() { + if [[ "$(id -u)" = "0" ]]; then + true + else + false + fi +} + +######################## +# Print OS metadata +# Arguments: +# $1 - Flag name +# Flags: +# --id - Distro ID +# --version - Distro version +# --branch - Distro branch +# --codename - Distro codename +# --name - Distro name +# --pretty-name - Distro pretty name +# Returns: +# String +######################### +get_os_metadata() { + local -r flag_name="${1:?missing flag}" + # Helper function + get_os_release_metadata() { + local -r env_name="${1:?missing environment variable name}" + ( + . /etc/os-release + echo "${!env_name}" + ) + } + case "$flag_name" in + --id) + get_os_release_metadata ID + ;; + --version) + get_os_release_metadata VERSION_ID + ;; + --branch) + get_os_release_metadata VERSION_ID | sed 's/\..*//' + ;; + --codename) + get_os_release_metadata VERSION_CODENAME + ;; + --name) + get_os_release_metadata NAME + ;; + --pretty-name) + get_os_release_metadata PRETTY_NAME + ;; + *) + error "Unknown flag ${flag_name}" + return 1 + ;; + esac +} + +######################## +# Get total memory available +# Arguments: +# None +# Returns: +# Memory in bytes +######################### +get_total_memory() { + echo $(($(grep MemTotal /proc/meminfo | awk '{print $2}') / 1024)) +} + +######################## +# Get machine size depending on specified memory +# Globals: +# None +# Arguments: +# None +# Flags: +# --memory - memory size (optional) +# Returns: +# Detected instance size +######################### +get_machine_size() { + local memory="" + # Validate arguments + while [[ "$#" -gt 0 ]]; do + case "$1" in + --memory) + shift + memory="${1:?missing memory}" + ;; + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + if [[ -z "$memory" ]]; then + debug "Memory was not specified, detecting available memory automatically" + memory="$(get_total_memory)" + fi + sanitized_memory=$(convert_to_mb "$memory") + if [[ "$sanitized_memory" -gt 26000 ]]; then + echo 2xlarge + elif [[ "$sanitized_memory" -gt 13000 ]]; then + echo xlarge + elif [[ "$sanitized_memory" -gt 6000 ]]; then + echo large + elif [[ "$sanitized_memory" -gt 3000 ]]; then + echo medium + elif [[ "$sanitized_memory" -gt 1500 ]]; then + echo small + else + echo micro + fi +} + +######################## +# Get machine size depending on specified memory +# Globals: +# None +# Arguments: +# $1 - memory size (optional) +# Returns: +# Detected instance size +######################### +get_supported_machine_sizes() { + echo micro small medium large xlarge 2xlarge +} + +######################## +# Convert memory size from string to amount of megabytes (i.e. 2G -> 2048) +# Globals: +# None +# Arguments: +# $1 - memory size +# Returns: +# Result of the conversion +######################### +convert_to_mb() { + local amount="${1:-}" + if [[ $amount =~ ^([0-9]+)(m|M|g|G) ]]; then + size="${BASH_REMATCH[1]}" + unit="${BASH_REMATCH[2]}" + if [[ "$unit" = "g" || "$unit" = "G" ]]; then + amount="$((size * 1024))" + else + amount="$size" + fi + fi + echo "$amount" +} + +######################### +# Redirects output to /dev/null if debug mode is disabled +# Globals: +# BITNAMI_DEBUG +# Arguments: +# $@ - Command to execute +# Returns: +# None +######################### +debug_execute() { + if is_boolean_yes "${BITNAMI_DEBUG:-false}"; then + "$@" + else + "$@" >/dev/null 2>&1 + fi +} + +######################## +# Retries a command a given number of times +# Arguments: +# $1 - cmd (as a string) +# $2 - max retries. Default: 12 +# $3 - sleep between retries (in seconds). Default: 5 +# Returns: +# Boolean +######################### +retry_while() { + local cmd="${1:?cmd is missing}" + local retries="${2:-12}" + local sleep_time="${3:-5}" + local return_value=1 + + read -r -a command <<<"$cmd" + for ((i = 1; i <= retries; i += 1)); do + "${command[@]}" && return_value=0 && break + sleep "$sleep_time" + done + return $return_value +} + +######################## +# Generate a random string +# Arguments: +# -t|--type - String type (ascii, alphanumeric, numeric), defaults to ascii +# -c|--count - Number of characters, defaults to 32 +# Arguments: +# None +# Returns: +# None +# Returns: +# String +######################### +generate_random_string() { + local type="ascii" + local count="32" + local filter + local result + # Validate arguments + while [[ "$#" -gt 0 ]]; do + case "$1" in + -t | --type) + shift + type="$1" + ;; + -c | --count) + shift + count="$1" + ;; + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + # Validate type + case "$type" in + ascii) + filter="[:print:]" + ;; + numeric) + filter="0-9" + ;; + alphanumeric) + filter="a-zA-Z0-9" + ;; + alphanumeric+special|special+alphanumeric) + # Limit variety of special characters, so there is a higher chance of containing more alphanumeric characters + # Special characters are harder to write, and it could impact the overall UX if most passwords are too complex + filter='a-zA-Z0-9:@.,/+!=' + ;; + *) + echo "Invalid type ${type}" >&2 + return 1 + ;; + esac + # Obtain count + 10 lines from /dev/urandom to ensure that the resulting string has the expected size + # Note there is a very small chance of strings starting with EOL character + # Therefore, the higher amount of lines read, this will happen less frequently + result="$(head -n "$((count + 10))" /dev/urandom | tr -dc "$filter" | head -c "$count")" + echo "$result" +} + +######################## +# Create md5 hash from a string +# Arguments: +# $1 - string +# Returns: +# md5 hash - string +######################### +generate_md5_hash() { + local -r str="${1:?missing input string}" + echo -n "$str" | md5sum | awk '{print $1}' +} + +######################## +# Create sha1 hash from a string +# Arguments: +# $1 - string +# $2 - algorithm - 1 (default), 224, 256, 384, 512 +# Returns: +# sha1 hash - string +######################### +generate_sha_hash() { + local -r str="${1:?missing input string}" + local -r algorithm="${2:-1}" + echo -n "$str" | "sha${algorithm}sum" | awk '{print $1}' +} + +######################## +# Converts a string to its hexadecimal representation +# Arguments: +# $1 - string +# Returns: +# hexadecimal representation of the string +######################### +convert_to_hex() { + local -r str=${1:?missing input string} + local -i iterator + local char + for ((iterator = 0; iterator < ${#str}; iterator++)); do + char=${str:iterator:1} + printf '%x' "'${char}" + done +} + +######################## +# Get boot time +# Globals: +# None +# Arguments: +# None +# Returns: +# Boot time metadata +######################### +get_boot_time() { + stat /proc --format=%Y +} + +######################## +# Get machine ID +# Globals: +# None +# Arguments: +# None +# Returns: +# Machine ID +######################### +get_machine_id() { + local machine_id + if [[ -f /etc/machine-id ]]; then + machine_id="$(cat /etc/machine-id)" + fi + if [[ -z "$machine_id" ]]; then + # Fallback to the boot-time, which will at least ensure a unique ID in the current session + machine_id="$(get_boot_time)" + fi + echo "$machine_id" +} + +######################## +# Get the root partition's disk device ID (e.g. /dev/sda1) +# Globals: +# None +# Arguments: +# None +# Returns: +# Root partition disk ID +######################### +get_disk_device_id() { + local device_id="" + if grep -q ^/dev /proc/mounts; then + device_id="$(grep ^/dev /proc/mounts | awk '$2 == "/" { print $1 }' | tail -1)" + fi + # If it could not be autodetected, fallback to /dev/sda1 as a default + if [[ -z "$device_id" || ! -b "$device_id" ]]; then + device_id="/dev/sda1" + fi + echo "$device_id" +} + +######################## +# Get the root disk device ID (e.g. /dev/sda) +# Globals: +# None +# Arguments: +# None +# Returns: +# Root disk ID +######################### +get_root_disk_device_id() { + get_disk_device_id | sed -E 's/p?[0-9]+$//' +} + +######################## +# Get the root disk size in bytes +# Globals: +# None +# Arguments: +# None +# Returns: +# Root disk size in bytes +######################### +get_root_disk_size() { + fdisk -l "$(get_root_disk_device_id)" | grep 'Disk.*bytes' | sed -E 's/.*, ([0-9]+) bytes,.*/\1/' || true +} + +######################## +# Run command as a specific user and group (optional) +# Arguments: +# $1 - USER(:GROUP) to switch to +# $2..$n - command to execute +# Returns: +# Exit code of the specified command +######################### +run_as_user() { + run_chroot "$@" +} + +######################## +# Execute command as a specific user and group (optional), +# replacing the current process image +# Arguments: +# $1 - USER(:GROUP) to switch to +# $2..$n - command to execute +# Returns: +# Exit code of the specified command +######################### +exec_as_user() { + run_chroot --replace-process "$@" +} + +######################## +# Run a command using chroot +# Arguments: +# $1 - USER(:GROUP) to switch to +# $2..$n - command to execute +# Flags: +# -r | --replace-process - Replace the current process image (optional) +# Returns: +# Exit code of the specified command +######################### +run_chroot() { + local userspec + local user + local homedir + local replace=false + local -r cwd="$(pwd)" + + # Parse and validate flags + while [[ "$#" -gt 0 ]]; do + case "$1" in + -r | --replace-process) + replace=true + ;; + --) + shift + break + ;; + -*) + stderr_print "unrecognized flag $1" + return 1 + ;; + *) + break + ;; + esac + shift + done + + # Parse and validate arguments + if [[ "$#" -lt 2 ]]; then + echo "expected at least 2 arguments" + return 1 + else + userspec=$1 + shift + + # userspec can optionally include the group, so we parse the user + user=$(echo "$userspec" | cut -d':' -f1) + fi + + if ! am_i_root; then + error "Could not switch to '${userspec}': Operation not permitted" + return 1 + fi + + # Get the HOME directory for the user to switch, as chroot does + # not properly update this env and some scripts rely on it + homedir=$(eval echo "~${user}") + if [[ ! -d $homedir ]]; then + homedir="${HOME:-/}" + fi + + # Obtaining value for "$@" indirectly in order to properly support shell parameter expansion + if [[ "$replace" = true ]]; then + exec chroot --userspec="$userspec" / bash -c "cd ${cwd}; export HOME=${homedir}; exec \"\$@\"" -- "$@" + else + chroot --userspec="$userspec" / bash -c "cd ${cwd}; export HOME=${homedir}; exec \"\$@\"" -- "$@" + fi +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libpersistence.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libpersistence.sh new file mode 100644 index 0000000000..18445e7d27 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libpersistence.sh @@ -0,0 +1,124 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Bitnami persistence library +# Used for bringing persistence capabilities to applications that don't have clear separation of data and logic + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/libfs.sh +. /opt/bitnami/scripts/libos.sh +. /opt/bitnami/scripts/liblog.sh +. /opt/bitnami/scripts/libversion.sh + +# Functions + +######################## +# Persist an application directory +# Globals: +# BITNAMI_ROOT_DIR +# BITNAMI_VOLUME_DIR +# Arguments: +# $1 - App folder name +# $2 - List of app files to persist +# Returns: +# true if all steps succeeded, false otherwise +######################### +persist_app() { + local -r app="${1:?missing app}" + local -a files_to_restore + read -r -a files_to_persist <<< "$(tr ',;:' ' ' <<< "$2")" + local -r install_dir="${BITNAMI_ROOT_DIR}/${app}" + local -r persist_dir="${BITNAMI_VOLUME_DIR}/${app}" + # Persist the individual files + if [[ "${#files_to_persist[@]}" -le 0 ]]; then + warn "No files are configured to be persisted" + return + fi + pushd "$install_dir" >/dev/null || exit + local file_to_persist_relative file_to_persist_destination file_to_persist_destination_folder + local -r tmp_file="/tmp/perms.acl" + for file_to_persist in "${files_to_persist[@]}"; do + if [[ ! -f "$file_to_persist" && ! -d "$file_to_persist" ]]; then + error "Cannot persist '${file_to_persist}' because it does not exist" + return 1 + fi + file_to_persist_relative="$(relativize "$file_to_persist" "$install_dir")" + file_to_persist_destination="${persist_dir}/${file_to_persist_relative}" + file_to_persist_destination_folder="$(dirname "$file_to_persist_destination")" + # Get original permissions for existing files, which will be applied later + # Exclude the root directory with 'sed', to avoid issues when copying the entirety of it to a volume + getfacl -R "$file_to_persist_relative" | sed -E '/# file: (\..+|[^.])/,$!d' > "$tmp_file" + # Copy directories to the volume + ensure_dir_exists "$file_to_persist_destination_folder" + cp -Lr --preserve=links "$file_to_persist_relative" "$file_to_persist_destination_folder" + # Restore permissions + pushd "$persist_dir" >/dev/null || exit + if am_i_root; then + setfacl --restore="$tmp_file" + else + # When running as non-root, don't change ownership + setfacl --restore=<(grep -E -v '^# (owner|group):' "$tmp_file") + fi + popd >/dev/null || exit + done + popd >/dev/null || exit + rm -f "$tmp_file" + # Install the persisted files into the installation directory, via symlinks + restore_persisted_app "$@" +} + +######################## +# Restore a persisted application directory +# Globals: +# BITNAMI_ROOT_DIR +# BITNAMI_VOLUME_DIR +# FORCE_MAJOR_UPGRADE +# Arguments: +# $1 - App folder name +# $2 - List of app files to restore +# Returns: +# true if all steps succeeded, false otherwise +######################### +restore_persisted_app() { + local -r app="${1:?missing app}" + local -a files_to_restore + read -r -a files_to_restore <<< "$(tr ',;:' ' ' <<< "$2")" + local -r install_dir="${BITNAMI_ROOT_DIR}/${app}" + local -r persist_dir="${BITNAMI_VOLUME_DIR}/${app}" + # Restore the individual persisted files + if [[ "${#files_to_restore[@]}" -le 0 ]]; then + warn "No persisted files are configured to be restored" + return + fi + local file_to_restore_relative file_to_restore_origin file_to_restore_destination + for file_to_restore in "${files_to_restore[@]}"; do + file_to_restore_relative="$(relativize "$file_to_restore" "$install_dir")" + # We use 'realpath --no-symlinks' to ensure that the case of '.' is covered and the directory is removed + file_to_restore_origin="$(realpath --no-symlinks "${install_dir}/${file_to_restore_relative}")" + file_to_restore_destination="$(realpath --no-symlinks "${persist_dir}/${file_to_restore_relative}")" + rm -rf "$file_to_restore_origin" + ln -sfn "$file_to_restore_destination" "$file_to_restore_origin" + done +} + +######################## +# Check if an application directory was already persisted +# Globals: +# BITNAMI_VOLUME_DIR +# Arguments: +# $1 - App folder name +# Returns: +# true if all steps succeeded, false otherwise +######################### +is_app_initialized() { + local -r app="${1:?missing app}" + local -r persist_dir="${BITNAMI_VOLUME_DIR}/${app}" + if ! is_mounted_dir_empty "$persist_dir"; then + true + else + false + fi +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libservice.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libservice.sh new file mode 100644 index 0000000000..1f9b33096b --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libservice.sh @@ -0,0 +1,496 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for managing services + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/libvalidations.sh +. /opt/bitnami/scripts/liblog.sh + +# Functions + +######################## +# Read the provided pid file and returns a PID +# Arguments: +# $1 - Pid file +# Returns: +# PID +######################### +get_pid_from_file() { + local pid_file="${1:?pid file is missing}" + + if [[ -f "$pid_file" ]]; then + if [[ -n "$(< "$pid_file")" ]] && [[ "$(< "$pid_file")" -gt 0 ]]; then + echo "$(< "$pid_file")" + fi + fi +} + +######################## +# Check if a provided PID corresponds to a running service +# Arguments: +# $1 - PID +# Returns: +# Boolean +######################### +is_service_running() { + local pid="${1:?pid is missing}" + + kill -0 "$pid" 2>/dev/null +} + +######################## +# Stop a service by sending a termination signal to its pid +# Arguments: +# $1 - Pid file +# $2 - Signal number (optional) +# Returns: +# None +######################### +stop_service_using_pid() { + local pid_file="${1:?pid file is missing}" + local signal="${2:-}" + local pid + + pid="$(get_pid_from_file "$pid_file")" + [[ -z "$pid" ]] || ! is_service_running "$pid" && return + + if [[ -n "$signal" ]]; then + kill "-${signal}" "$pid" + else + kill "$pid" + fi + + local counter=10 + while [[ "$counter" -ne 0 ]] && is_service_running "$pid"; do + sleep 1 + counter=$((counter - 1)) + done +} + +######################## +# Start cron daemon +# Arguments: +# None +# Returns: +# true if started correctly, false otherwise +######################### +cron_start() { + if [[ -x "/usr/sbin/cron" ]]; then + /usr/sbin/cron + elif [[ -x "/usr/sbin/crond" ]]; then + /usr/sbin/crond + else + false + fi +} + +######################## +# Generate a cron configuration file for a given service +# Arguments: +# $1 - Service name +# $2 - Command +# Flags: +# --run-as - User to run as (default: root) +# --schedule - Cron schedule configuration (default: * * * * *) +# Returns: +# None +######################### +generate_cron_conf() { + local service_name="${1:?service name is missing}" + local cmd="${2:?command is missing}" + local run_as="root" + local schedule="* * * * *" + local clean="true" + + # Parse optional CLI flags + shift 2 + while [[ "$#" -gt 0 ]]; do + case "$1" in + --run-as) + shift + run_as="$1" + ;; + --schedule) + shift + schedule="$1" + ;; + --no-clean) + clean="false" + ;; + *) + echo "Invalid command line flag ${1}" >&2 + return 1 + ;; + esac + shift + done + + mkdir -p /etc/cron.d + if "$clean"; then + cat > "/etc/cron.d/${service_name}" <> /etc/cron.d/"$service_name" + fi +} + +######################## +# Remove a cron configuration file for a given service +# Arguments: +# $1 - Service name +# Returns: +# None +######################### +remove_cron_conf() { + local service_name="${1:?service name is missing}" + local cron_conf_dir="/etc/monit/conf.d" + rm -f "${cron_conf_dir}/${service_name}" +} + +######################## +# Generate a monit configuration file for a given service +# Arguments: +# $1 - Service name +# $2 - Pid file +# $3 - Start command +# $4 - Stop command +# Flags: +# --disable - Whether to disable the monit configuration +# Returns: +# None +######################### +generate_monit_conf() { + local service_name="${1:?service name is missing}" + local pid_file="${2:?pid file is missing}" + local start_command="${3:?start command is missing}" + local stop_command="${4:?stop command is missing}" + local monit_conf_dir="/etc/monit/conf.d" + local disabled="no" + + # Parse optional CLI flags + shift 4 + while [[ "$#" -gt 0 ]]; do + case "$1" in + --disable) + disabled="yes" + ;; + *) + echo "Invalid command line flag ${1}" >&2 + return 1 + ;; + esac + shift + done + + is_boolean_yes "$disabled" && conf_suffix=".disabled" + mkdir -p "$monit_conf_dir" + cat > "${monit_conf_dir}/${service_name}.conf${conf_suffix:-}" <&2 + return 1 + ;; + esac + shift + done + + mkdir -p "$logrotate_conf_dir" + cat < "${logrotate_conf_dir}/${service_name}" +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +${log_path} { + ${period} + rotate ${rotations} + dateext + compress + copytruncate + missingok +$(indent "$extra" 2) +} +EOF +} + +######################## +# Remove a logrotate configuration file +# Arguments: +# $1 - Service name +# Returns: +# None +######################### +remove_logrotate_conf() { + local service_name="${1:?service name is missing}" + local logrotate_conf_dir="/etc/logrotate.d" + rm -f "${logrotate_conf_dir}/${service_name}" +} + +######################## +# Generate a Systemd configuration file +# Arguments: +# $1 - Service name +# Flags: +# --custom-service-content - Custom content to add to the [service] block +# --environment - Environment variable to define (multiple --environment options may be passed) +# --environment-file - Text file with environment variables (multiple --environment-file options may be passed) +# --exec-start - Start command (required) +# --exec-start-pre - Pre-start command (optional) +# --exec-start-post - Post-start command (optional) +# --exec-stop - Stop command (optional) +# --exec-reload - Reload command (optional) +# --group - System group to start the service with +# --name - Service full name (e.g. Apache HTTP Server, defaults to $1) +# --restart - When to restart the Systemd service after being stopped (defaults to always) +# --pid-file - Service PID file +# --standard-output - File where to print stdout output +# --standard-error - File where to print stderr output +# --success-exit-status - Exit code that indicates a successful shutdown +# --type - Systemd unit type (defaults to forking) +# --user - System user to start the service with +# --working-directory - Working directory at which to start the service +# Returns: +# None +######################### +generate_systemd_conf() { + local -r service_name="${1:?service name is missing}" + local -r systemd_units_dir="/etc/systemd/system" + local -r service_file="${systemd_units_dir}/bitnami.${service_name}.service" + # Default values + local name="$service_name" + local type="forking" + local user="" + local group="" + local environment="" + local environment_file="" + local exec_start="" + local exec_start_pre="" + local exec_start_post="" + local exec_stop="" + local exec_reload="" + local restart="always" + local pid_file="" + local standard_output="journal" + local standard_error="" + local limits_content="" + local success_exit_status="" + local custom_service_content="" + local working_directory="" + # Parse CLI flags + shift + while [[ "$#" -gt 0 ]]; do + case "$1" in + --name \ + | --type \ + | --user \ + | --group \ + | --exec-start \ + | --exec-stop \ + | --exec-reload \ + | --restart \ + | --pid-file \ + | --standard-output \ + | --standard-error \ + | --success-exit-status \ + | --custom-service-content \ + | --working-directory \ + ) + var_name="$(echo "$1" | sed -e "s/^--//" -e "s/-/_/g")" + shift + declare "$var_name"="${1:?"${var_name} value is missing"}" + ;; + --limit-*) + [[ -n "$limits_content" ]] && limits_content+=$'\n' + var_name="${1//--limit-}" + shift + limits_content+="Limit${var_name^^}=${1:?"--limit-${var_name} value is missing"}" + ;; + --exec-start-pre) + shift + [[ -n "$exec_start_pre" ]] && exec_start_pre+=$'\n' + exec_start_pre+="ExecStartPre=${1:?"--exec-start-pre value is missing"}" + ;; + --exec-start-post) + shift + [[ -n "$exec_start_post" ]] && exec_start_post+=$'\n' + exec_start_post+="ExecStartPost=${1:?"--exec-start-post value is missing"}" + ;; + --environment) + shift + # It is possible to add multiple environment lines + [[ -n "$environment" ]] && environment+=$'\n' + environment+="Environment=${1:?"--environment value is missing"}" + ;; + --environment-file) + shift + # It is possible to add multiple environment-file lines + [[ -n "$environment_file" ]] && environment_file+=$'\n' + environment_file+="EnvironmentFile=${1:?"--environment-file value is missing"}" + ;; + *) + echo "Invalid command line flag ${1}" >&2 + return 1 + ;; + esac + shift + done + # Validate inputs + local error="no" + if [[ -z "$exec_start" ]]; then + error "The --exec-start option is required" + error="yes" + fi + if [[ "$error" != "no" ]]; then + return 1 + fi + # Generate the Systemd unit + cat > "$service_file" <> "$service_file" <<< "WorkingDirectory=${working_directory}" + fi + if [[ -n "$exec_start_pre" ]]; then + # This variable may contain multiple ExecStartPre= directives + cat >> "$service_file" <<< "$exec_start_pre" + fi + if [[ -n "$exec_start" ]]; then + cat >> "$service_file" <<< "ExecStart=${exec_start}" + fi + if [[ -n "$exec_start_post" ]]; then + # This variable may contain multiple ExecStartPost= directives + cat >> "$service_file" <<< "$exec_start_post" + fi + # Optional stop and reload commands + if [[ -n "$exec_stop" ]]; then + cat >> "$service_file" <<< "ExecStop=${exec_stop}" + fi + if [[ -n "$exec_reload" ]]; then + cat >> "$service_file" <<< "ExecReload=${exec_reload}" + fi + # User and group + if [[ -n "$user" ]]; then + cat >> "$service_file" <<< "User=${user}" + fi + if [[ -n "$group" ]]; then + cat >> "$service_file" <<< "Group=${group}" + fi + # PID file allows to determine if the main process is running properly (for Restart=always) + if [[ -n "$pid_file" ]]; then + cat >> "$service_file" <<< "PIDFile=${pid_file}" + fi + if [[ -n "$restart" ]]; then + cat >> "$service_file" <<< "Restart=${restart}" + fi + # Environment flags + if [[ -n "$environment" ]]; then + # This variable may contain multiple Environment= directives + cat >> "$service_file" <<< "$environment" + fi + if [[ -n "$environment_file" ]]; then + # This variable may contain multiple EnvironmentFile= directives + cat >> "$service_file" <<< "$environment_file" + fi + # Logging + if [[ -n "$standard_output" ]]; then + cat >> "$service_file" <<< "StandardOutput=${standard_output}" + fi + if [[ -n "$standard_error" ]]; then + cat >> "$service_file" <<< "StandardError=${standard_error}" + fi + if [[ -n "$custom_service_content" ]]; then + # This variable may contain multiple miscellaneous directives + cat >> "$service_file" <<< "$custom_service_content" + fi + if [[ -n "$success_exit_status" ]]; then + cat >> "$service_file" <> "$service_file" <> "$service_file" <> "$service_file" <= 0 )); then + true + else + false + fi +} + +######################## +# Check if the provided argument is a boolean or is the string 'yes/true' +# Arguments: +# $1 - Value to check +# Returns: +# Boolean +######################### +is_boolean_yes() { + local -r bool="${1:-}" + # comparison is performed without regard to the case of alphabetic characters + shopt -s nocasematch + if [[ "$bool" = 1 || "$bool" =~ ^(yes|true)$ ]]; then + true + else + false + fi +} + +######################## +# Check if the provided argument is a boolean yes/no value +# Arguments: +# $1 - Value to check +# Returns: +# Boolean +######################### +is_yes_no_value() { + local -r bool="${1:-}" + if [[ "$bool" =~ ^(yes|no)$ ]]; then + true + else + false + fi +} + +######################## +# Check if the provided argument is a boolean true/false value +# Arguments: +# $1 - Value to check +# Returns: +# Boolean +######################### +is_true_false_value() { + local -r bool="${1:-}" + if [[ "$bool" =~ ^(true|false)$ ]]; then + true + else + false + fi +} + +######################## +# Check if the provided argument is a boolean 1/0 value +# Arguments: +# $1 - Value to check +# Returns: +# Boolean +######################### +is_1_0_value() { + local -r bool="${1:-}" + if [[ "$bool" =~ ^[10]$ ]]; then + true + else + false + fi +} + +######################## +# Check if the provided argument is an empty string or not defined +# Arguments: +# $1 - Value to check +# Returns: +# Boolean +######################### +is_empty_value() { + local -r val="${1:-}" + if [[ -z "$val" ]]; then + true + else + false + fi +} + +######################## +# Validate if the provided argument is a valid port +# Arguments: +# $1 - Port to validate +# Returns: +# Boolean and error message +######################### +validate_port() { + local value + local unprivileged=0 + + # Parse flags + while [[ "$#" -gt 0 ]]; do + case "$1" in + -unprivileged) + unprivileged=1 + ;; + --) + shift + break + ;; + -*) + stderr_print "unrecognized flag $1" + return 1 + ;; + *) + break + ;; + esac + shift + done + + if [[ "$#" -gt 1 ]]; then + echo "too many arguments provided" + return 2 + elif [[ "$#" -eq 0 ]]; then + stderr_print "missing port argument" + return 1 + else + value=$1 + fi + + if [[ -z "$value" ]]; then + echo "the value is empty" + return 1 + else + if ! is_int "$value"; then + echo "value is not an integer" + return 2 + elif [[ "$value" -lt 0 ]]; then + echo "negative value provided" + return 2 + elif [[ "$value" -gt 65535 ]]; then + echo "requested port is greater than 65535" + return 2 + elif [[ "$unprivileged" = 1 && "$value" -lt 1024 ]]; then + echo "privileged port requested" + return 3 + fi + fi +} + +######################## +# Validate if the provided argument is a valid IPv6 address +# Arguments: +# $1 - IP to validate +# Returns: +# Boolean +######################### +validate_ipv6() { + local ip="${1:?ip is missing}" + local stat=1 + local full_address_regex='^([0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}$' + local short_address_regex='^((([0-9a-fA-F]{1,4}:){0,6}[0-9a-fA-F]{1,4}){0,6}::(([0-9a-fA-F]{1,4}:){0,6}[0-9a-fA-F]{1,4}){0,6})$' + + if [[ $ip =~ $full_address_regex || $ip =~ $short_address_regex || $ip == "::" ]]; then + stat=0 + fi + return $stat +} + +######################## +# Validate if the provided argument is a valid IPv4 address +# Arguments: +# $1 - IP to validate +# Returns: +# Boolean +######################### +validate_ipv4() { + local ip="${1:?ip is missing}" + local stat=1 + + if [[ $ip =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then + read -r -a ip_array <<< "$(tr '.' ' ' <<< "$ip")" + [[ ${ip_array[0]} -le 255 && ${ip_array[1]} -le 255 \ + && ${ip_array[2]} -le 255 && ${ip_array[3]} -le 255 ]] + stat=$? + fi + return $stat +} + +######################## +# Validate if the provided argument is a valid IPv4 or IPv6 address +# Arguments: +# $1 - IP to validate +# Returns: +# Boolean +######################### +validate_ip() { + local ip="${1:?ip is missing}" + local stat=1 + + if validate_ipv4 "$ip"; then + stat=0 + else + stat=$(validate_ipv6 "$ip") + fi + return $stat +} + +######################## +# Validate a string format +# Arguments: +# $1 - String to validate +# Returns: +# Boolean +######################### +validate_string() { + local string + local min_length=-1 + local max_length=-1 + + # Parse flags + while [ "$#" -gt 0 ]; do + case "$1" in + -min-length) + shift + min_length=${1:-} + ;; + -max-length) + shift + max_length=${1:-} + ;; + --) + shift + break + ;; + -*) + stderr_print "unrecognized flag $1" + return 1 + ;; + *) + break + ;; + esac + shift + done + + if [ "$#" -gt 1 ]; then + stderr_print "too many arguments provided" + return 2 + elif [ "$#" -eq 0 ]; then + stderr_print "missing string" + return 1 + else + string=$1 + fi + + if [[ "$min_length" -ge 0 ]] && [[ "${#string}" -lt "$min_length" ]]; then + echo "string length is less than $min_length" + return 1 + fi + if [[ "$max_length" -ge 0 ]] && [[ "${#string}" -gt "$max_length" ]]; then + echo "string length is great than $max_length" + return 1 + fi +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libversion.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libversion.sh new file mode 100644 index 0000000000..f0d5a5cd33 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libversion.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Library for managing versions strings + +# shellcheck disable=SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/liblog.sh + +# Functions +######################## +# Gets semantic version +# Arguments: +# $1 - version: string to extract major.minor.patch +# $2 - section: 1 to extract major, 2 to extract minor, 3 to extract patch +# Returns: +# array with the major, minor and release +######################### +get_sematic_version () { + local version="${1:?version is required}" + local section="${2:?section is required}" + local -a version_sections + + #Regex to parse versions: x.y.z + local -r regex='([0-9]+)(\.([0-9]+)(\.([0-9]+))?)?' + + if [[ "$version" =~ $regex ]]; then + local i=1 + local j=1 + local n=${#BASH_REMATCH[*]} + + while [[ $i -lt $n ]]; do + if [[ -n "${BASH_REMATCH[$i]}" ]] && [[ "${BASH_REMATCH[$i]:0:1}" != '.' ]]; then + version_sections[j]="${BASH_REMATCH[$i]}" + ((j++)) + fi + ((i++)) + done + + local number_regex='^[0-9]+$' + if [[ "$section" =~ $number_regex ]] && (( section > 0 )) && (( section <= 3 )); then + echo "${version_sections[$section]}" + return + else + stderr_print "Section allowed values are: 1, 2, and 3" + return 1 + fi + fi +} diff --git a/container/bitnami/prebuildfs/opt/bitnami/scripts/libwebserver.sh b/container/bitnami/prebuildfs/opt/bitnami/scripts/libwebserver.sh new file mode 100644 index 0000000000..acb84fc233 --- /dev/null +++ b/container/bitnami/prebuildfs/opt/bitnami/scripts/libwebserver.sh @@ -0,0 +1,476 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Bitnami web server handler library + +# shellcheck disable=SC1090,SC1091 + +# Load generic libraries +. /opt/bitnami/scripts/liblog.sh + +######################## +# Execute a command (or list of commands) with the web server environment and library loaded +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_execute() { + local -r web_server="${1:?missing web server}" + shift + # Run program in sub-shell to avoid web server environment getting loaded when not necessary + ( + . "/opt/bitnami/scripts/lib${web_server}.sh" + . "/opt/bitnami/scripts/${web_server}-env.sh" + "$@" + ) +} + +######################## +# Prints the list of enabled web servers +# Globals: +# None +# Arguments: +# None +# Returns: +# None +######################### +web_server_list() { + local -r -a supported_web_servers=(apache nginx) + local -a existing_web_servers=() + for web_server in "${supported_web_servers[@]}"; do + [[ -f "/opt/bitnami/scripts/${web_server}-env.sh" ]] && existing_web_servers+=("$web_server") + done + echo "${existing_web_servers[@]:-}" +} + +######################## +# Prints the currently-enabled web server type (only one, in order of preference) +# Globals: +# None +# Arguments: +# None +# Returns: +# None +######################### +web_server_type() { + local -a web_servers + read -r -a web_servers <<< "$(web_server_list)" + echo "${web_servers[0]:-}" +} + +######################## +# Validate that a supported web server is configured +# Globals: +# None +# Arguments: +# None +# Returns: +# None +######################### +web_server_validate() { + local error_code=0 + local supported_web_servers=("apache" "nginx") + + # Auxiliary functions + print_validation_error() { + error "$1" + error_code=1 + } + + if [[ -z "$(web_server_type)" || ! " ${supported_web_servers[*]} " == *" $(web_server_type) "* ]]; then + print_validation_error "Could not detect any supported web servers. It must be one of: ${supported_web_servers[*]}" + elif ! web_server_execute "$(web_server_type)" type -t "is_$(web_server_type)_running" >/dev/null; then + print_validation_error "Could not load the $(web_server_type) web server library from /opt/bitnami/scripts. Check that it exists and is readable." + fi + + return "$error_code" +} + +######################## +# Check whether the web server is running +# Globals: +# * +# Arguments: +# None +# Returns: +# true if the web server is running, false otherwise +######################### +is_web_server_running() { + "is_$(web_server_type)_running" +} + +######################## +# Start web server +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_start() { + info "Starting $(web_server_type) in background" + if [[ "${BITNAMI_SERVICE_MANAGER:-}" = "systemd" ]]; then + systemctl start "bitnami.$(web_server_type).service" + else + "${BITNAMI_ROOT_DIR}/scripts/$(web_server_type)/start.sh" + fi +} + +######################## +# Stop web server +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_stop() { + info "Stopping $(web_server_type)" + if [[ "${BITNAMI_SERVICE_MANAGER:-}" = "systemd" ]]; then + systemctl stop "bitnami.$(web_server_type).service" + else + "${BITNAMI_ROOT_DIR}/scripts/$(web_server_type)/stop.sh" + fi +} + +######################## +# Restart web server +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_restart() { + info "Restarting $(web_server_type)" + if [[ "${BITNAMI_SERVICE_MANAGER:-}" = "systemd" ]]; then + systemctl restart "bitnami.$(web_server_type).service" + else + "${BITNAMI_ROOT_DIR}/scripts/$(web_server_type)/restart.sh" + fi +} + +######################## +# Reload web server +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_reload() { + if [[ "${BITNAMI_SERVICE_MANAGER:-}" = "systemd" ]]; then + systemctl reload "bitnami.$(web_server_type).service" + else + "${BITNAMI_ROOT_DIR}/scripts/$(web_server_type)/reload.sh" + fi +} + +######################## +# Ensure a web server application configuration exists (i.e. Apache virtual host format or NGINX server block) +# It serves as a wrapper for the specific web server function +# Globals: +# * +# Arguments: +# $1 - App name +# Flags: +# --type - Application type, which has an effect on which configuration template to use +# --hosts - Host listen addresses +# --server-name - Server name +# --server-aliases - Server aliases +# --allow-remote-connections - Whether to allow remote connections or to require local connections +# --disable - Whether to render server configurations with a .disabled prefix +# --disable-http - Whether to render the app's HTTP server configuration with a .disabled prefix +# --disable-https - Whether to render the app's HTTPS server configuration with a .disabled prefix +# --http-port - HTTP port number +# --https-port - HTTPS port number +# --document-root - Path to document root directory +# Apache-specific flags: +# --apache-additional-configuration - Additional vhost configuration (no default) +# --apache-additional-http-configuration - Additional HTTP vhost configuration (no default) +# --apache-additional-https-configuration - Additional HTTPS vhost configuration (no default) +# --apache-before-vhost-configuration - Configuration to add before the directive (no default) +# --apache-allow-override - Whether to allow .htaccess files (only allowed when --move-htaccess is set to 'no' and type is not defined) +# --apache-extra-directory-configuration - Extra configuration for the document root directory +# --apache-proxy-address - Address where to proxy requests +# --apache-proxy-configuration - Extra configuration for the proxy +# --apache-proxy-http-configuration - Extra configuration for the proxy HTTP vhost +# --apache-proxy-https-configuration - Extra configuration for the proxy HTTPS vhost +# --apache-move-htaccess - Move .htaccess files to a common place so they can be loaded during Apache startup (only allowed when type is not defined) +# NGINX-specific flags: +# --nginx-additional-configuration - Additional server block configuration (no default) +# --nginx-external-configuration - Configuration external to server block (no default) +# Returns: +# true if the configuration was enabled, false otherwise +######################## +ensure_web_server_app_configuration_exists() { + local app="${1:?missing app}" + shift + local -a apache_args nginx_args web_servers args_var + apache_args=("$app") + nginx_args=("$app") + # Validate arguments + while [[ "$#" -gt 0 ]]; do + case "$1" in + # Common flags + --disable \ + | --disable-http \ + | --disable-https \ + ) + apache_args+=("$1") + nginx_args+=("$1") + ;; + --hosts \ + | --server-name \ + | --server-aliases \ + | --type \ + | --allow-remote-connections \ + | --http-port \ + | --https-port \ + | --document-root \ + ) + apache_args+=("$1" "${2:?missing value}") + nginx_args+=("$1" "${2:?missing value}") + shift + ;; + + # Specific Apache flags + --apache-additional-configuration \ + | --apache-additional-http-configuration \ + | --apache-additional-https-configuration \ + | --apache-before-vhost-configuration \ + | --apache-allow-override \ + | --apache-extra-directory-configuration \ + | --apache-proxy-address \ + | --apache-proxy-configuration \ + | --apache-proxy-http-configuration \ + | --apache-proxy-https-configuration \ + | --apache-move-htaccess \ + ) + apache_args+=("${1//apache-/}" "${2:?missing value}") + shift + ;; + + # Specific NGINX flags + --nginx-additional-configuration \ + | --nginx-external-configuration) + nginx_args+=("${1//nginx-/}" "${2:?missing value}") + shift + ;; + + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + read -r -a web_servers <<< "$(web_server_list)" + for web_server in "${web_servers[@]}"; do + args_var="${web_server}_args[@]" + web_server_execute "$web_server" "ensure_${web_server}_app_configuration_exists" "${!args_var}" + done +} + +######################## +# Ensure a web server application configuration does not exist anymore (i.e. Apache virtual host format or NGINX server block) +# It serves as a wrapper for the specific web server function +# Globals: +# * +# Arguments: +# $1 - App name +# Returns: +# true if the configuration was disabled, false otherwise +######################## +ensure_web_server_app_configuration_not_exists() { + local app="${1:?missing app}" + local -a web_servers + read -r -a web_servers <<< "$(web_server_list)" + for web_server in "${web_servers[@]}"; do + web_server_execute "$web_server" "ensure_${web_server}_app_configuration_not_exists" "$app" + done +} + +######################## +# Ensure the web server loads the configuration for an application in a URL prefix +# It serves as a wrapper for the specific web server function +# Globals: +# * +# Arguments: +# $1 - App name +# Flags: +# --allow-remote-connections - Whether to allow remote connections or to require local connections +# --document-root - Path to document root directory +# --prefix - URL prefix from where it will be accessible (i.e. /myapp) +# --type - Application type, which has an effect on what configuration template will be used +# Apache-specific flags: +# --apache-additional-configuration - Additional vhost configuration (no default) +# --apache-allow-override - Whether to allow .htaccess files (only allowed when --move-htaccess is set to 'no') +# --apache-extra-directory-configuration - Extra configuration for the document root directory +# --apache-move-htaccess - Move .htaccess files to a common place so they can be loaded during Apache startup +# NGINX-specific flags: +# --nginx-additional-configuration - Additional server block configuration (no default) +# Returns: +# true if the configuration was enabled, false otherwise +######################## +ensure_web_server_prefix_configuration_exists() { + local app="${1:?missing app}" + shift + local -a apache_args nginx_args web_servers args_var + apache_args=("$app") + nginx_args=("$app") + # Validate arguments + while [[ "$#" -gt 0 ]]; do + case "$1" in + # Common flags + --allow-remote-connections \ + | --document-root \ + | --prefix \ + | --type \ + ) + apache_args+=("$1" "${2:?missing value}") + nginx_args+=("$1" "${2:?missing value}") + shift + ;; + + # Specific Apache flags + --apache-additional-configuration \ + | --apache-allow-override \ + | --apache-extra-directory-configuration \ + | --apache-move-htaccess \ + ) + apache_args+=("${1//apache-/}" "$2") + shift + ;; + + # Specific NGINX flags + --nginx-additional-configuration) + nginx_args+=("${1//nginx-/}" "$2") + shift + ;; + + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + read -r -a web_servers <<< "$(web_server_list)" + for web_server in "${web_servers[@]}"; do + args_var="${web_server}_args[@]" + web_server_execute "$web_server" "ensure_${web_server}_prefix_configuration_exists" "${!args_var}" + done +} + +######################## +# Ensure a web server application configuration is updated with the runtime configuration (i.e. ports) +# It serves as a wrapper for the specific web server function +# Globals: +# * +# Arguments: +# $1 - App name +# Flags: +# --hosts - Host listen addresses +# --server-name - Server name +# --server-aliases - Server aliases +# --enable-http - Enable HTTP app configuration (if not enabled already) +# --enable-https - Enable HTTPS app configuration (if not enabled already) +# --disable-http - Disable HTTP app configuration (if not disabled already) +# --disable-https - Disable HTTPS app configuration (if not disabled already) +# --http-port - HTTP port number +# --https-port - HTTPS port number +# Returns: +# true if the configuration was updated, false otherwise +######################## +web_server_update_app_configuration() { + local app="${1:?missing app}" + shift + local -a args web_servers + args=("$app") + # Validate arguments + while [[ "$#" -gt 0 ]]; do + case "$1" in + # Common flags + --enable-http \ + | --enable-https \ + | --disable-http \ + | --disable-https \ + ) + args+=("$1") + ;; + --hosts \ + | --server-name \ + | --server-aliases \ + | --http-port \ + | --https-port \ + ) + args+=("$1" "${2:?missing value}") + shift + ;; + + *) + echo "Invalid command line flag $1" >&2 + return 1 + ;; + esac + shift + done + read -r -a web_servers <<< "$(web_server_list)" + for web_server in "${web_servers[@]}"; do + web_server_execute "$web_server" "${web_server}_update_app_configuration" "${args[@]}" + done +} + +######################## +# Enable loading page, which shows users that the initialization process is not yet completed +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_enable_loading_page() { + ensure_web_server_app_configuration_exists "__loading" --hosts "_default_" \ + --apache-additional-configuration " +# Show a HTTP 503 Service Unavailable page by default +RedirectMatch 503 ^/$ +# Show index.html if server is answering with 404 Not Found or 503 Service Unavailable status codes +ErrorDocument 404 /index.html +ErrorDocument 503 /index.html" \ + --nginx-additional-configuration " +# Show a HTTP 503 Service Unavailable page by default +location / { + return 503; +} +# Show index.html if server is answering with 404 Not Found or 503 Service Unavailable status codes +error_page 404 @installing; +error_page 503 @installing; +location @installing { + rewrite ^(.*)$ /index.html break; +}" + web_server_reload +} + +######################## +# Enable loading page, which shows users that the initialization process is not yet completed +# Globals: +# * +# Arguments: +# None +# Returns: +# None +######################### +web_server_disable_install_page() { + ensure_web_server_app_configuration_not_exists "__loading" + web_server_reload +} diff --git a/container/bitnami/prebuildfs/usr/sbin/install_packages b/container/bitnami/prebuildfs/usr/sbin/install_packages new file mode 100644 index 0000000000..ccce248b2d --- /dev/null +++ b/container/bitnami/prebuildfs/usr/sbin/install_packages @@ -0,0 +1,27 @@ +#!/bin/sh +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +set -eu + +n=0 +max=2 +export DEBIAN_FRONTEND=noninteractive + +until [ $n -gt $max ]; do + set +e + ( + apt-get update -qq && + apt-get install -y --no-install-recommends "$@" + ) + CODE=$? + set -e + if [ $CODE -eq 0 ]; then + break + fi + if [ $n -eq $max ]; then + exit $CODE + fi + echo "apt failed, retrying" + n=$(($n + 1)) +done +apt-get clean && rm -rf /var/lib/apt/lists /var/cache/apt/archives diff --git a/container/bitnami/prebuildfs/usr/sbin/run-script b/container/bitnami/prebuildfs/usr/sbin/run-script new file mode 100644 index 0000000000..0e07c9038d --- /dev/null +++ b/container/bitnami/prebuildfs/usr/sbin/run-script @@ -0,0 +1,24 @@ +#!/bin/sh +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +set -u + +if [ $# -eq 0 ]; then + >&2 echo "No arguments provided" + exit 1 +fi + +script=$1 +exit_code="${2:-96}" +fail_if_not_present="${3:-n}" + +if test -f "$script"; then + sh $script + + if [ $? -ne 0 ]; then + exit $((exit_code)) + fi +elif [ "$fail_if_not_present" = "y" ]; then + >&2 echo "script not found: $script" + exit 127 +fi diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/java/entrypoint.sh b/container/bitnami/rootfs/opt/bitnami/scripts/java/entrypoint.sh new file mode 100644 index 0000000000..38802fc0bf --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/java/entrypoint.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +# shellcheck disable=SC1091 + +set -o errexit +set -o nounset +set -o pipefail +# set -o xtrace # Uncomment this line for debugging purposes + +# Load libraries +. /opt/bitnami/scripts/libbitnami.sh +. /opt/bitnami/scripts/liblog.sh + +if [[ "$OS_FLAVOUR" =~ photon && "$APP_VERSION" =~ ^1.8 ]]; then + # Option --module-path is not supported by JAVA 1.8 since modules were added in version 1.9 + unset JAVA_TOOL_OPTIONS +fi + +print_welcome_page + +echo "" +exec "$@" diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/java/postunpack.sh b/container/bitnami/rootfs/opt/bitnami/scripts/java/postunpack.sh new file mode 100644 index 0000000000..c15af32e2f --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/java/postunpack.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +# shellcheck disable=SC1091 + +set -o errexit +set -o nounset +set -o pipefail +# set -o xtrace # Uncomment this line for debugging purposes + +# Load libraries +. /opt/bitnami/scripts/libfile.sh +. /opt/bitnami/scripts/liblog.sh + +# +# Java post-unpack operations +# + +# Override default files in the Java security directory. This is used for +# custom base images (with custom CA certificates or block lists is used) + +if [[ -n "${JAVA_EXTRA_SECURITY_DIR:-}" ]] && ! is_dir_empty "$JAVA_EXTRA_SECURITY_DIR"; then + info "Adding custom CAs to the Java security folder" + cp -Lr "${JAVA_EXTRA_SECURITY_DIR}/." /opt/bitnami/java/lib/security +fi diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/kafka-env.sh b/container/bitnami/rootfs/opt/bitnami/scripts/kafka-env.sh new file mode 100644 index 0000000000..e1621b93e8 --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/kafka-env.sh @@ -0,0 +1,121 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Environment configuration for kafka + +# The values for all environment variables will be set in the below order of precedence +# 1. Custom environment variables defined below after Bitnami defaults +# 2. Constants defined in this file (environment variables with no default), i.e. BITNAMI_ROOT_DIR +# 3. Environment variables overridden via external files using *_FILE variables (see below) +# 4. Environment variables set externally (i.e. current Bash context/Dockerfile/userdata) + +# Load logging library +# shellcheck disable=SC1090,SC1091 +. /opt/bitnami/scripts/liblog.sh + +export BITNAMI_ROOT_DIR="/opt/bitnami" +export BITNAMI_VOLUME_DIR="/bitnami" + +# Logging configuration +export MODULE="${MODULE:-kafka}" +export BITNAMI_DEBUG="${BITNAMI_DEBUG:-false}" + +# By setting an environment variable matching *_FILE to a file path, the prefixed environment +# variable will be overridden with the value specified in that file +kafka_env_vars=( + KAFKA_MOUNTED_CONF_DIR + KAFKA_INTER_BROKER_USER + KAFKA_INTER_BROKER_PASSWORD + KAFKA_CONTROLLER_USER + KAFKA_CONTROLLER_PASSWORD + KAFKA_CERTIFICATE_PASSWORD + KAFKA_TLS_TRUSTSTORE_FILE + KAFKA_TLS_TYPE + KAFKA_TLS_CLIENT_AUTH + KAFKA_OPTS + KAFKA_CFG_SASL_ENABLED_MECHANISMS + KAFKA_KRAFT_CLUSTER_ID + KAFKA_SKIP_KRAFT_STORAGE_INIT + KAFKA_CLIENT_LISTENER_NAME + KAFKA_ZOOKEEPER_PROTOCOL + KAFKA_ZOOKEEPER_PASSWORD + KAFKA_ZOOKEEPER_USER + KAFKA_ZOOKEEPER_TLS_KEYSTORE_PASSWORD + KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_PASSWORD + KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE + KAFKA_ZOOKEEPER_TLS_VERIFY_HOSTNAME + KAFKA_ZOOKEEPER_TLS_TYPE + KAFKA_CLIENT_USERS + KAFKA_CLIENT_PASSWORDS + KAFKA_HEAP_OPTS + JAVA_TOOL_OPTIONS +) +for env_var in "${kafka_env_vars[@]}"; do + file_env_var="${env_var}_FILE" + if [[ -n "${!file_env_var:-}" ]]; then + if [[ -r "${!file_env_var:-}" ]]; then + export "${env_var}=$(< "${!file_env_var}")" + unset "${file_env_var}" + else + warn "Skipping export of '${env_var}'. '${!file_env_var:-}' is not readable." + fi + fi +done +unset kafka_env_vars + +# Paths +export KAFKA_BASE_DIR="${BITNAMI_ROOT_DIR}/kafka" +export KAFKA_VOLUME_DIR="/bitnami/kafka" +export KAFKA_DATA_DIR="${KAFKA_VOLUME_DIR}/data" +export KAFKA_CONF_DIR="${KAFKA_BASE_DIR}/config" +export KAFKA_CONF_FILE="${KAFKA_CONF_DIR}/server.properties" +export KAFKA_MOUNTED_CONF_DIR="${KAFKA_MOUNTED_CONF_DIR:-${KAFKA_VOLUME_DIR}/config}" +export KAFKA_CERTS_DIR="${KAFKA_CONF_DIR}/certs" +export KAFKA_INITSCRIPTS_DIR="/docker-entrypoint-initdb.d" +export KAFKA_LOG_DIR="${KAFKA_BASE_DIR}/logs" +export KAFKA_HOME="$KAFKA_BASE_DIR" +export PATH="${KAFKA_BASE_DIR}/bin:${BITNAMI_ROOT_DIR}/java/bin:${PATH}" + +# System users (when running with a privileged user) +export KAFKA_DAEMON_USER="kafka" +export KAFKA_DAEMON_GROUP="kafka" + +# Kafka runtime settings +export KAFKA_INTER_BROKER_USER="${KAFKA_INTER_BROKER_USER:-user}" +export KAFKA_INTER_BROKER_PASSWORD="${KAFKA_INTER_BROKER_PASSWORD:-bitnami}" +export KAFKA_CONTROLLER_USER="${KAFKA_CONTROLLER_USER:-controller_user}" +export KAFKA_CONTROLLER_PASSWORD="${KAFKA_CONTROLLER_PASSWORD:-bitnami}" +export KAFKA_CERTIFICATE_PASSWORD="${KAFKA_CERTIFICATE_PASSWORD:-}" +export KAFKA_TLS_TRUSTSTORE_FILE="${KAFKA_TLS_TRUSTSTORE_FILE:-}" +export KAFKA_TLS_TYPE="${KAFKA_TLS_TYPE:-JKS}" +export KAFKA_TLS_CLIENT_AUTH="${KAFKA_TLS_CLIENT_AUTH:-required}" +export KAFKA_OPTS="${KAFKA_OPTS:-}" + +# Kafka configuration overrides +export KAFKA_CFG_SASL_ENABLED_MECHANISMS="${KAFKA_CFG_SASL_ENABLED_MECHANISMS:-PLAIN,SCRAM-SHA-256,SCRAM-SHA-512}" +export KAFKA_KRAFT_CLUSTER_ID="${KAFKA_KRAFT_CLUSTER_ID:-}" +export KAFKA_SKIP_KRAFT_STORAGE_INIT="${KAFKA_SKIP_KRAFT_STORAGE_INIT:-false}" +export KAFKA_CLIENT_LISTENER_NAME="${KAFKA_CLIENT_LISTENER_NAME:-}" + +# ZooKeeper connection settings +export KAFKA_ZOOKEEPER_PROTOCOL="${KAFKA_ZOOKEEPER_PROTOCOL:-PLAINTEXT}" +export KAFKA_ZOOKEEPER_PASSWORD="${KAFKA_ZOOKEEPER_PASSWORD:-}" +export KAFKA_ZOOKEEPER_USER="${KAFKA_ZOOKEEPER_USER:-}" +export KAFKA_ZOOKEEPER_TLS_KEYSTORE_PASSWORD="${KAFKA_ZOOKEEPER_TLS_KEYSTORE_PASSWORD:-}" +export KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_PASSWORD="${KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_PASSWORD:-}" +export KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE="${KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE:-}" +export KAFKA_ZOOKEEPER_TLS_VERIFY_HOSTNAME="${KAFKA_ZOOKEEPER_TLS_VERIFY_HOSTNAME:-true}" +export KAFKA_ZOOKEEPER_TLS_TYPE="${KAFKA_ZOOKEEPER_TLS_TYPE:-JKS}" + +# Authentication +export KAFKA_CLIENT_USERS="${KAFKA_CLIENT_USERS:-user}" +export KAFKA_CLIENT_PASSWORDS="${KAFKA_CLIENT_PASSWORDS:-bitnami}" + +# Java settings +export KAFKA_HEAP_OPTS="${KAFKA_HEAP_OPTS:--Xmx1024m -Xms1024m}" + +# Java settings +export JAVA_TOOL_OPTIONS="${JAVA_TOOL_OPTIONS:-}" + +# Custom environment variables may be defined below diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/kafka/entrypoint.sh b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/entrypoint.sh new file mode 100644 index 0000000000..d7413bcfc4 --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/entrypoint.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +# shellcheck disable=SC1091 + +set -o errexit +set -o nounset +set -o pipefail +# set -o xtrace # Uncomment this line for debugging purposes + +# Load libraries +. /opt/bitnami/scripts/liblog.sh +. /opt/bitnami/scripts/libbitnami.sh +. /opt/bitnami/scripts/libkafka.sh + +# Load Kafka environment variables +. /opt/bitnami/scripts/kafka-env.sh + +print_welcome_page + +if [[ "$*" = *"/opt/bitnami/scripts/kafka/run.sh"* || "$*" = *"/run.sh"* ]]; then + info "** Starting Kafka setup **" + /opt/bitnami/scripts/kafka/setup.sh + info "** Kafka setup finished! **" +fi + +echo "" +exec "$@" diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh new file mode 100644 index 0000000000..673c84e721 --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/postunpack.sh @@ -0,0 +1,46 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +# shellcheck disable=SC1091 + +set -o errexit +set -o nounset +set -o pipefail +# set -o xtrace # Uncomment this line for debugging purposes + +# Load libraries +. /opt/bitnami/scripts/libkafka.sh +. /opt/bitnami/scripts/libfs.sh + +# Load Kafka environment variables +. /opt/bitnami/scripts/kafka-env.sh + +# Move server.properties from configtmp to config +# Temporary solution until kafka tarball places server.properties into config +if [[ -d "${KAFKA_BASE_DIR}/configtmp" ]]; then + mv "${KAFKA_BASE_DIR}/configtmp"/* "$KAFKA_CONF_DIR" + rmdir "${KAFKA_BASE_DIR}/configtmp" +fi +[[ -d "${KAFKA_BASE_DIR}/conf" ]] && rmdir "${KAFKA_BASE_DIR}/conf" + +# Ensure directories used by Kafka exist and have proper ownership and permissions +for dir in "$KAFKA_LOG_DIR" "$KAFKA_CONF_DIR" "$KAFKA_MOUNTED_CONF_DIR" "$KAFKA_VOLUME_DIR" "$KAFKA_DATA_DIR" "$KAFKA_INITSCRIPTS_DIR"; do + ensure_dir_exists "$dir" +done +chmod -R g+rwX "$KAFKA_BASE_DIR" "$KAFKA_VOLUME_DIR" "$KAFKA_DATA_DIR" "$KAFKA_INITSCRIPTS_DIR" + +# Move the original server.properties, so users can skip initialization logic by mounting their own server.properties directly instead of using the MOUNTED_CONF_DIR +mv "${KAFKA_CONF_DIR}/server.properties" "${KAFKA_CONF_DIR}/server.properties.original" + +# Disable logging to stdout and garbage collection +# Source: https://logging.apache.org/log4j/log4j-2.4/manual/appenders.html +#replace_in_file "${KAFKA_BASE_DIR}/bin/kafka-server-start.sh" " [-]loggc" " " +#replace_in_file "${KAFKA_CONF_DIR}/log4j.properties" "DailyRollingFileAppender" "ConsoleAppender" + +# Disable the default console logger in favour of KafkaAppender (which provides the exact output) +#echo "log4j.appender.stdout.Threshold=OFF" >>"${KAFKA_CONF_DIR}/log4j.properties" + +# Remove invalid parameters for ConsoleAppender +#remove_in_file "${KAFKA_CONF_DIR}/log4j.properties" "DatePattern" +#remove_in_file "${KAFKA_CONF_DIR}/log4j.properties" "Appender.File" diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/kafka/run.sh b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/run.sh new file mode 100644 index 0000000000..76d4380aff --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/run.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +# shellcheck disable=SC1091 + +set -o errexit +set -o nounset +set -o pipefail +# set -o xtrace # Uncomment this line for debugging purposes + +# Load libraries +. /opt/bitnami/scripts/libkafka.sh +. /opt/bitnami/scripts/libos.sh + +# Load Kafka environment variables +. /opt/bitnami/scripts/kafka-env.sh + +if [[ -f "${KAFKA_CONF_DIR}/kafka_jaas.conf" ]]; then + export KAFKA_OPTS="${KAFKA_OPTS:-} -Djava.security.auth.login.config=${KAFKA_CONF_DIR}/kafka_jaas.conf" +fi + +cmd="$KAFKA_HOME/bin/kafka-server-start.sh" +args=("$KAFKA_CONF_FILE") +! is_empty_value "${KAFKA_EXTRA_FLAGS:-}" && args=("${args[@]}" "${KAFKA_EXTRA_FLAGS[@]}") + +info "** Starting Kafka **" +if am_i_root; then + exec_as_user "$KAFKA_DAEMON_USER" "$cmd" "${args[@]}" "$@" +else + exec "$cmd" "${args[@]}" "$@" +fi diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/kafka/setup.sh b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/setup.sh new file mode 100644 index 0000000000..5195f71910 --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/kafka/setup.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 + +# shellcheck disable=SC1091 + +set -o errexit +set -o nounset +set -o pipefail +# set -o xtrace # Uncomment this line for debugging purposes + +# Load libraries +. /opt/bitnami/scripts/libfs.sh +. /opt/bitnami/scripts/libos.sh +. /opt/bitnami/scripts/libkafka.sh + +# Load Kafka environment variables +. /opt/bitnami/scripts/kafka-env.sh + +# Map Kafka environment variables +kafka_create_alias_environment_variables + +# Dinamically set node.id/broker.id/controller.quorum.voters if the _COMMAND environment variable is set +kafka_dynamic_environment_variables + +# Set the default tuststore locations before validation +kafka_configure_default_truststore_locations +# Ensure Kafka user and group exist when running as 'root' +am_i_root && ensure_user_exists "$KAFKA_DAEMON_USER" --group "$KAFKA_DAEMON_GROUP" +# Ensure directories used by Kafka exist and have proper ownership and permissions +for dir in "$KAFKA_LOG_DIR" "$KAFKA_CONF_DIR" "$KAFKA_MOUNTED_CONF_DIR" "$KAFKA_VOLUME_DIR" "$KAFKA_DATA_DIR"; do + if am_i_root; then + ensure_dir_exists "$dir" "$KAFKA_DAEMON_USER" "$KAFKA_DAEMON_GROUP" + else + ensure_dir_exists "$dir" + fi +done + +# Kafka validation, skipped if server.properties was mounted at either $KAFKA_MOUNTED_CONF_DIR or $KAFKA_CONF_DIR +[[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/server.properties" && ! -f "$KAFKA_CONF_FILE" ]] && kafka_validate +# Kafka initialization, skipped if server.properties was mounted at $KAFKA_CONF_DIR +[[ ! -f "$KAFKA_CONF_FILE" ]] && kafka_initialize + +# Initialise KRaft metadata storage if process.roles configured +if grep -q "^process.roles=" "$KAFKA_CONF_FILE" && ! is_boolean_yes "$KAFKA_SKIP_KRAFT_STORAGE_INIT" ; then + kafka_kraft_storage_initialize +fi +# Configure Zookeeper SCRAM users +if is_boolean_yes "${KAFKA_ZOOKEEPER_BOOTSTRAP_SCRAM_USERS:-}"; then + kafka_zookeeper_create_sasl_scram_users +fi +# KRaft controllers may get stuck starting when the controller quorum voters are changed. +# Workaround: Remove quorum-state file when scaling up/down controllers (Waiting proposal KIP-853) +# https://cwiki.apache.org/confluence/display/KAFKA/KIP-853%3A+KRaft+Voter+Changes +if [[ -f "${KAFKA_DATA_DIR}/__cluster_metadata-0/quorum-state" ]] && grep -q "^controller.quorum.voters=" "$KAFKA_CONF_FILE" && kafka_kraft_quorum_voters_changed; then + warn "Detected inconsitences between controller.quorum.voters and quorum-state, removing it..." + rm -f "${KAFKA_DATA_DIR}/__cluster_metadata-0/quorum-state" +fi +# Ensure custom initialization scripts are executed +kafka_custom_init_scripts diff --git a/container/bitnami/rootfs/opt/bitnami/scripts/libkafka.sh b/container/bitnami/rootfs/opt/bitnami/scripts/libkafka.sh new file mode 100644 index 0000000000..df2459b035 --- /dev/null +++ b/container/bitnami/rootfs/opt/bitnami/scripts/libkafka.sh @@ -0,0 +1,1176 @@ +#!/bin/bash +# Copyright Broadcom, Inc. All Rights Reserved. +# SPDX-License-Identifier: APACHE-2.0 +# +# Bitnami Kafka library + +# shellcheck disable=SC1090,SC1091 + +# Load Generic Libraries +. /opt/bitnami/scripts/libfile.sh +. /opt/bitnami/scripts/libfs.sh +. /opt/bitnami/scripts/liblog.sh +. /opt/bitnami/scripts/libos.sh +. /opt/bitnami/scripts/libvalidations.sh +. /opt/bitnami/scripts/libservice.sh + +# Functions + +######################## +# Set a configuration setting value to a file +# Globals: +# None +# Arguments: +# $1 - file +# $2 - key +# $3 - values (array) +# Returns: +# None +######################### +kafka_common_conf_set() { + local file="${1:?missing file}" + local key="${2:?missing key}" + shift + shift + local values=("$@") + + if [[ "${#values[@]}" -eq 0 ]]; then + stderr_print "missing value" + return 1 + elif [[ "${#values[@]}" -ne 1 ]]; then + for i in "${!values[@]}"; do + kafka_common_conf_set "$file" "${key[$i]}" "${values[$i]}" + done + else + value="${values[0]}" + # Check if the value was set before + if grep -q "^[#\\s]*$key\s*=.*" "$file"; then + # Update the existing key + replace_in_file "$file" "^[#\\s]*${key}\s*=.*" "${key}=${value}" false + else + # Add a new key + printf '\n%s=%s' "$key" "$value" >>"$file" + fi + fi +} + +######################## +# Returns true if at least one listener is configured using SSL +# Globals: +# KAFKA_CFG_LISTENERS +# KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP +# Arguments: +# None +# Returns: +# true/false +######################### +kafka_has_ssl_listener(){ + if ! is_empty_value "${KAFKA_CFG_LISTENERS:-}"; then + if is_empty_value "${KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP:-}"; then + if [[ "$KAFKA_CFG_LISTENERS" =~ SSL: || "$KAFKA_CFG_LISTENERS" =~ SASL_SSL: ]]; then + return + fi + else + read -r -a protocol_maps <<<"$(tr ',' ' ' <<<"$KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP")" + for protocol_map in "${protocol_maps[@]}"; do + read -r -a map <<<"$(tr ':' ' ' <<<"$protocol_map")" + # Obtain the listener and protocol from protocol map string, e.g. CONTROLLER:PLAINTEXT + listener="${map[0]}" + protocol="${map[1]}" + if [[ "$protocol" = "SSL" || "$protocol" = "SASL_SSL" ]]; then + if [[ "$KAFKA_CFG_LISTENERS" =~ $listener ]]; then + return + fi + fi + done + fi + fi + return 1 +} + +######################## +# Returns true if at least one listener is configured using SASL +# Globals: +# KAFKA_CFG_LISTENERS +# KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP +# Arguments: +# None +# Returns: +# true/false +######################### +kafka_has_sasl_listener(){ + if ! is_empty_value "${KAFKA_CFG_LISTENERS:-}"; then + if is_empty_value "${KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP:-}"; then + if [[ "$KAFKA_CFG_LISTENERS" =~ SASL_PLAINTEXT: ]] || [[ "$KAFKA_CFG_LISTENERS" =~ SASL_SSL: ]]; then + return + fi + else + read -r -a protocol_maps <<<"$(tr ',' ' ' <<<"$KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP")" + for protocol_map in "${protocol_maps[@]}"; do + read -r -a map <<<"$(tr ':' ' ' <<<"$protocol_map")" + # Obtain the listener and protocol from protocol map string, e.g. CONTROLLER:PLAINTEXT + listener="${map[0]}" + protocol="${map[1]}" + if [[ "$protocol" = "SASL_PLAINTEXT" || "$protocol" = "SASL_SSL" ]]; then + if [[ "$KAFKA_CFG_LISTENERS" =~ $listener ]]; then + return + fi + fi + done + fi + fi + return 1 +} + +######################## +# Returns true if at least one listener is configured using plaintext +# Globals: +# KAFKA_CFG_LISTENERS +# KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP +# Arguments: +# None +# Returns: +# true/false +######################### +kafka_has_plaintext_listener(){ + if ! is_empty_value "${KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP:-}"; then + read -r -a protocol_maps <<<"$(tr ',' ' ' <<<"$KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP")" + for protocol_map in "${protocol_maps[@]}"; do + read -r -a map <<<"$(tr ':' ' ' <<<"$protocol_map")" + # Obtain the listener and protocol from protocol map string, e.g. CONTROLLER:PLAINTEXT + listener="${map[0]}" + protocol="${map[1]}" + if [[ "$protocol" = "PLAINTEXT" ]]; then + if is_empty_value "${KAFKA_CFG_LISTENERS:-}" || [[ "$KAFKA_CFG_LISTENERS" =~ $listener ]]; then + return + fi + fi + done + else + if is_empty_value "${KAFKA_CFG_LISTENERS:-}" || [[ "$KAFKA_CFG_LISTENERS" =~ PLAINTEXT: ]]; then + return + fi + fi + return 1 +} + +######################## +# Backwards compatibility measure to configure the TLS truststore locations +# Globals: +# KAFKA_CONF_FILE +# Arguments: +# None +# Returns: +# None +######################### +kafka_configure_default_truststore_locations() { + # Backwards compatibility measure to allow custom truststore locations but at the same time not disrupt + # the UX that the previous version of the containers and the helm chart have. + # Context: The chart and containers by default assumed that the truststore location was KAFKA_CERTS_DIR/kafka.truststore.jks or KAFKA_MOUNTED_CONF_DIR/certs/kafka.truststore.jks. + # Because of this, we could not use custom certificates in different locations (use case: A custom base image that already has a truststore). Changing the logic to allow custom + # locations implied major changes in the current user experience (which only required to mount certificates at the assumed location). In order to maintain this compatibility we need + # use this logic that sets the KAFKA_TLS_*_FILE variables to the previously assumed locations in case it is not set + + # Kafka truststore + if kafka_has_ssl_listener && is_empty_value "${KAFKA_TLS_TRUSTSTORE_FILE:-}"; then + local kafka_truststore_filename="kafka.truststore.jks" + [[ "$KAFKA_TLS_TYPE" = "PEM" ]] && kafka_truststore_filename="kafka.truststore.pem" + if [[ -f "${KAFKA_CERTS_DIR}/${kafka_truststore_filename}" ]]; then + # Mounted in /opt/bitnami/kafka/conf/certs + export KAFKA_TLS_TRUSTSTORE_FILE="${KAFKA_CERTS_DIR}/${kafka_truststore_filename}" + else + # Mounted in /bitnami/kafka/conf/certs + export KAFKA_TLS_TRUSTSTORE_FILE="${KAFKA_MOUNTED_CONF_DIR}/certs/${kafka_truststore_filename}" + fi + fi + # Zookeeper truststore + if [[ "${KAFKA_ZOOKEEPER_PROTOCOL:-}" =~ SSL ]] && is_empty_value "${KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE:-}"; then + local zk_truststore_filename="zookeeper.truststore.jks" + [[ "$KAFKA_ZOOKEEPER_TLS_TYPE" = "PEM" ]] && zk_truststore_filename="zookeeper.truststore.pem" + if [[ -f "${KAFKA_CERTS_DIR}/${zk_truststore_filename}" ]]; then + # Mounted in /opt/bitnami/kafka/conf/certs + export KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE="${KAFKA_CERTS_DIR}/${zk_truststore_filename}" + else + # Mounted in /bitnami/kafka/conf/certs + export KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE="${KAFKA_MOUNTED_CONF_DIR}/certs/${zk_truststore_filename}" + fi + fi +} + +######################## +# Set a configuration setting value to server.properties +# Globals: +# KAFKA_CONF_FILE +# Arguments: +# $1 - key +# $2 - values (array) +# Returns: +# None +######################### +kafka_server_conf_set() { + kafka_common_conf_set "$KAFKA_CONF_FILE" "$@" +} + +######################## +# Set a configuration setting value to producer.properties and consumer.properties +# Globals: +# KAFKA_CONF_DIR +# Arguments: +# $1 - key +# $2 - values (array) +# Returns: +# None +######################### +kafka_producer_consumer_conf_set() { + kafka_common_conf_set "$KAFKA_CONF_DIR/producer.properties" "$@" + kafka_common_conf_set "$KAFKA_CONF_DIR/consumer.properties" "$@" +} + +######################## +# Create alias for environment variable, so both can be used +# Globals: +# None +# Arguments: +# $1 - Alias environment variable name +# $2 - Original environment variable name +# Returns: +# None +######################### +kafka_declare_alias_env() { + local -r alias="${1:?missing environment variable alias}" + local -r original="${2:?missing original environment variable}" + if printenv "${original}" >/dev/null; then + export "$alias"="${!original:-}" + fi +} + +######################## +# Map Kafka legacy environment variables to the new names +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_create_alias_environment_variables() { + suffixes=( + "ADVERTISED_LISTENERS" + "BROKER_ID" + "NODE_ID" + "CONTROLLER_QUORUM_VOTERS" + "PROCESS_ROLES" + "DEFAULT_REPLICATION_FACTOR" + "DELETE_TOPIC_ENABLE" + "INTER_BROKER_LISTENER_NAME" + "LISTENERS" + "LISTENER_SECURITY_PROTOCOL_MAP" + "LOG_DIRS" + "LOG_FLUSH_INTERVAL_MESSAGES" + "LOG_FLUSH_INTERVAL_MS" + "LOG_MESSAGE_FORMAT_VERSION" + "LOG_RETENTION_BYTES" + "LOG_RETENTION_CHECK_INTERVALS_MS" + "LOG_RETENTION_HOURS" + "LOG_SEGMENT_BYTES" + "MESSAGE_MAX_BYTES" + "NUM_IO_THREADS" + "NUM_NETWORK_THREADS" + "NUM_PARTITIONS" + "NUM_RECOVERY_THREADS_PER_DATA_DIR" + "OFFSETS_TOPIC_REPLICATION_FACTOR" + "SOCKET_RECEIVE_BUFFER_BYTES" + "SOCKET_REQUEST_MAX_BYTES" + "SOCKET_SEND_BUFFER_BYTES" + "SSL_ENDPOINT_IDENTIFICATION_ALGORITHM" + "TRANSACTION_STATE_LOG_MIN_ISR" + "TRANSACTION_STATE_LOG_REPLICATION_FACTOR" + "ZOOKEEPER_CONNECT" + "ZOOKEEPER_CONNECTION_TIMEOUT_MS" + ) + kafka_declare_alias_env "KAFKA_CFG_LOG_DIRS" "KAFKA_LOGS_DIRS" + kafka_declare_alias_env "KAFKA_CFG_LOG_SEGMENT_BYTES" "KAFKA_SEGMENT_BYTES" + kafka_declare_alias_env "KAFKA_CFG_MESSAGE_MAX_BYTES" "KAFKA_MAX_MESSAGE_BYTES" + kafka_declare_alias_env "KAFKA_CFG_ZOOKEEPER_CONNECTION_TIMEOUT_MS" "KAFKA_ZOOKEEPER_CONNECT_TIMEOUT_MS" + kafka_declare_alias_env "KAFKA_CFG_AUTO_CREATE_TOPICS_ENABLE" "KAFKA_AUTO_CREATE_TOPICS_ENABLE" + kafka_declare_alias_env "KAFKA_CLIENT_USERS" "KAFKA_BROKER_USER" + kafka_declare_alias_env "KAFKA_CLIENT_PASSWORDS" "KAFKA_BROKER_PASSWORD" + kafka_declare_alias_env "KAFKA_CLIENT_LISTENER_NAME" "KAFKA_CLIENT_LISTENER" + for s in "${suffixes[@]}"; do + kafka_declare_alias_env "KAFKA_CFG_${s}" "KAFKA_${s}" + done +} + +######################## +# Validate settings in KAFKA_* env vars +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_validate() { + debug "Validating settings in KAFKA_* env vars..." + local error_code=0 + + # Auxiliary functions + print_validation_error() { + error "$1" + error_code=1 + } + check_multi_value() { + if [[ " ${2} " != *" ${!1} "* ]]; then + print_validation_error "The allowed values for ${1} are: ${2}" + fi + } + # If process.roles configured, check its values are valid and perform additional checks for each + check_kraft_process_roles() { + read -r -a roles_list <<<"$(tr ',;' ' ' <<<"$KAFKA_CFG_PROCESS_ROLES")" + for role in "${roles_list[@]}"; do + case "$role" in + broker) ;; + controller) + if is_empty_value "${KAFKA_CFG_CONTROLLER_LISTENER_NAMES:-}"; then + print_validation_error "Role 'controller' enabled but environment variable KAFKA_CFG_CONTROLLER_LISTENER_NAMES was not provided." + fi + if is_empty_value "${KAFKA_CFG_LISTENERS:-}" || [[ ! "$KAFKA_CFG_LISTENERS" =~ ${KAFKA_CFG_CONTROLLER_LISTENER_NAMES} ]]; then + print_validation_error "Role 'controller' enabled but listener ${KAFKA_CFG_CONTROLLER_LISTENER_NAMES} not found in KAFKA_CFG_LISTENERS." + fi + ;; + *) + print_validation_error "Invalid KRaft process role '$role'. Supported roles are 'broker,controller'" + ;; + esac + done + } + # Check all listeners are using a unique and valid port + check_listener_ports(){ + check_allowed_port() { + local port="${1:?missing port variable}" + local -a validate_port_args=() + ! am_i_root && validate_port_args+=("-unprivileged") + validate_port_args+=("$port") + if ! err=$(validate_port "${validate_port_args[@]}"); then + print_validation_error "An invalid port ${port} was specified in the environment variable KAFKA_CFG_LISTENERS: ${err}." + fi + } + + read -r -a listeners <<<"$(tr ',' ' ' <<<"${KAFKA_CFG_LISTENERS:-}")" + local -a ports=() + for listener in "${listeners[@]}"; do + read -r -a arr <<<"$(tr ':' ' ' <<<"$listener")" + # Obtain the port from listener string, e.g. PLAINTEXT://:9092 + port="${arr[2]}" + check_allowed_port "$port" + ports+=("$port") + done + # Check each listener is using an unique port + local -a unique_ports=() + read -r -a unique_ports <<< "$(echo "${ports[@]}" | tr ' ' '\n' | sort -u | tr '\n' ' ')" + if [[ "${#ports[@]}" != "${#unique_ports[@]}" ]]; then + print_validation_error "There are listeners bound to the same port" + fi + } + check_listener_protocols(){ + local -r allowed_protocols=("PLAINTEXT" "SASL_PLAINTEXT" "SASL_SSL" "SSL") + read -r -a protocol_maps <<<"$(tr ',' ' ' <<<"$KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP")" + for protocol_map in "${protocol_maps[@]}"; do + read -r -a map <<<"$(tr ':' ' ' <<<"$protocol_map")" + # Obtain the listener and protocol from protocol map string, e.g. CONTROLLER:PLAINTEXT + listener="${map[0]}" + protocol="${map[1]}" + # Check protocol in allowed list + if [[ ! "${allowed_protocols[*]}" =~ $protocol ]]; then + print_validation_error "Authentication protocol ${protocol} is not supported!" + fi + # If inter-broker listener configured with SASL, ensure KAFKA_CFG_SASL_MECHANISM_INTER_BROKER_PROTOCOL is set + if [[ "$listener" = "${KAFKA_CFG_INTER_BROKER_LISTENER_NAME:-INTERNAL}" ]]; then + if [[ "$protocol" = "SASL_PLAINTEXT" ]] || [[ "$protocol" = "SASL_SSL" ]]; then + if is_empty_value "${KAFKA_CFG_SASL_MECHANISM_INTER_BROKER_PROTOCOL:-}"; then + print_validation_error "When using SASL for inter broker comunication the mechanism should be provided using KAFKA_CFG_SASL_MECHANISM_INTER_BROKER_PROTOCOL" + fi + if is_empty_value "${KAFKA_INTER_BROKER_USER:-}" || is_empty_value "${KAFKA_INTER_BROKER_PASSWORD:-}"; then + print_validation_error "In order to configure SASL authentication for Kafka inter-broker communications, you must provide the SASL credentials. Set the environment variables KAFKA_INTER_BROKER_USER and KAFKA_INTER_BROKER_PASSWORD to configure the credentials for SASL authentication with between brokers." + fi + fi + # If controller listener configured with SASL, ensure KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL is set + elif [[ "${KAFKA_CFG_CONTROLLER_LISTENER_NAMES:-CONTROLLER}" =~ $listener ]]; then + if [[ "$protocol" = "SASL_PLAINTEXT" ]] || [[ "$protocol" = "SASL_SSL" ]]; then + if is_empty_value "${KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL:-}"; then + print_validation_error "When using SASL for controller comunication the mechanism should be provided at KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" + elif [[ "$KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" =~ SCRAM ]]; then + warn "KRaft controller listener may not support SCRAM-SHA-256/SCRAM-SHA-512 mechanisms. If facing any issues, we recommend switching to PLAIN mechanism. More information at: https://issues.apache.org/jira/browse/KAFKA-15513" + fi + if is_empty_value "${KAFKA_CONTROLLER_USER:-}" || is_empty_value "${KAFKA_CONTROLLER_PASSWORD:-}"; then + print_validation_error "In order to configure SASL authentication for Kafka control plane communications, you must provide the SASL credentials. Set the environment variables KAFKA_CONTROLLER_USER and KAFKA_CONTROLLER_PASSWORD to configure the credentials for SASL authentication with between controllers." + fi + fi + else + if [[ "$protocol" = "SASL_PLAINTEXT" ]] || [[ "$protocol" = "SASL_SSL" ]]; then + if is_empty_value "${KAFKA_CLIENT_USERS:-}" || is_empty_value "${KAFKA_CLIENT_PASSWORDS:-}"; then + print_validation_error "In order to configure SASL authentication for Kafka, you must provide the SASL credentials. Set the environment variables KAFKA_CLIENT_USERS and KAFKA_CLIENT_PASSWORDS to configure the credentials for SASL authentication with clients." + fi + fi + + fi + done + } + + if is_empty_value "${KAFKA_CFG_PROCESS_ROLES:-}" && is_empty_value "${KAFKA_CFG_ZOOKEEPER_CONNECT:-}"; then + print_validation_error "Kafka haven't been configured to work in either Raft or Zookeper mode. Please make sure at least one of the modes is configured." + fi + # Check KRaft mode + if ! is_empty_value "${KAFKA_CFG_PROCESS_ROLES:-}"; then + # Only allow Zookeeper configuration if migration mode is enabled + if ! is_empty_value "${KAFKA_CFG_ZOOKEEPER_CONNECT:-}" && + { is_empty_value "${KAFKA_CFG_ZOOKEEPER_METADATA_MIGRATION_ENABLE:-}" || ! is_boolean_yes "$KAFKA_CFG_ZOOKEEPER_METADATA_MIGRATION_ENABLE"; }; then + print_validation_error "Both KRaft mode and Zookeeper modes are configured, but KAFKA_CFG_ZOOKEEPER_METADATA_MIGRATION_ENABLE is not enabled" + fi + if is_empty_value "${KAFKA_CFG_NODE_ID:-}"; then + print_validation_error "KRaft mode requires an unique node.id, please set the environment variable KAFKA_CFG_NODE_ID" + fi + if is_empty_value "${KAFKA_CFG_CONTROLLER_QUORUM_VOTERS:-}"; then + print_validation_error "KRaft mode requires KAFKA_CFG_CONTROLLER_QUORUM_VOTERS to be set" + fi + check_kraft_process_roles + fi + # Check Zookeeper mode + if ! is_empty_value "${KAFKA_CFG_ZOOKEEPER_CONNECT:-}"; then + # If SSL/SASL_SSL protocol configured, check certificates are provided + if [[ "$KAFKA_ZOOKEEPER_PROTOCOL" =~ SSL ]]; then + if [[ "$KAFKA_ZOOKEEPER_TLS_TYPE" = "JKS" ]]; then + # Fail if truststore is not provided + if [[ ! -f "$KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE" ]]; then + print_validation_error "In order to configure the TLS encryption for Zookeeper with JKS certs you must mount your zookeeper.truststore.jks cert to the ${KAFKA_MOUNTED_CONF_DIR}/certs directory." + fi + # Warn if keystore is not provided, only required if Zookeper mTLS is enabled (ZOO_TLS_CLIENT_AUTH) + if [[ ! -f "${KAFKA_CERTS_DIR}/zookeeper.keystore.jks" ]] && [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/certs/zookeeper.keystore.jks" ]]; then + warn "In order to configure the mTLS for Zookeeper with JKS certs you must mount your zookeeper.keystore.jks cert to the ${KAFKA_MOUNTED_CONF_DIR}/certs directory." + fi + elif [[ "$KAFKA_ZOOKEEPER_TLS_TYPE" = "PEM" ]]; then + # Fail if CA / validation cert is not provided + if [[ ! -f "$KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE" ]]; then + print_validation_error "In order to configure the TLS encryption for Zookeeper with PEM certs you must mount your zookeeper.truststore.pem cert to the ${KAFKA_MOUNTED_CONF_DIR}/certs directory." + fi + # Warn if node key or cert are not provided, only required if Zookeper mTLS is enabled (ZOO_TLS_CLIENT_AUTH) + if { [[ ! -f "${KAFKA_CERTS_DIR}/zookeeper.keystore.pem" ]] || [[ ! -f "${KAFKA_CERTS_DIR}/zookeeper.keystore.key" ]]; } && + { [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/certs/zookeeper.keystore.pem" ]] || [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/certs/zookeeper.keystore.key" ]]; }; then + warn "In order to configure the mTLS for Zookeeper with PEM certs you must mount your zookeeper.keystore.pem cert and zookeeper.keystore.key key to the ${KAFKA_MOUNTED_CONF_DIR}/certs directory." + fi + fi + fi + # If SASL/SASL_SSL protocol configured, check certificates are provided + if [[ "$KAFKA_ZOOKEEPER_PROTOCOL" =~ SASL ]]; then + if is_empty_value "${KAFKA_ZOOKEEPER_USER:-}" || is_empty_value "${KAFKA_ZOOKEEPER_PASSWORD:-}"; then + print_validation_error "In order to configure SASL authentication for Kafka, you must provide the SASL credentials. Set the environment variables KAFKA_ZOOKEEPER_USER and KAFKA_ZOOKEEPER_PASSWORD, to configure the credentials for SASL authentication with Zookeeper." + fi + fi + # If using plaintext protocol, check it is explicitly allowed + if [[ "$KAFKA_ZOOKEEPER_PROTOCOL" = "PLAINTEXT" ]]; then + warn "The KAFKA_ZOOKEEPER_PROTOCOL environment variable does not configure SASL and/or SSL, this setting is not recommended for production environments." + fi + fi + # Check listener ports are unique and allowed + check_listener_ports + # Check listeners are mapped to a valid security protocol + check_listener_protocols + # Warn users if plaintext listeners are configured + if kafka_has_plaintext_listener; then + warn "Kafka has been configured with a PLAINTEXT listener, this setting is not recommended for production environments." + fi + # If SSL/SASL_SSL listeners configured, check certificates are provided + if kafka_has_ssl_listener; then + if [[ "$KAFKA_TLS_TYPE" = "JKS" ]] && + { [[ ! -f "${KAFKA_CERTS_DIR}/kafka.keystore.jks" ]] || [[ ! -f "$KAFKA_TLS_TRUSTSTORE_FILE" ]]; } && + { [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/certs/kafka.keystore.jks" ]] || [[ ! -f "$KAFKA_TLS_TRUSTSTORE_FILE" ]]; }; then + print_validation_error "In order to configure the TLS encryption for Kafka with JKS certs you must mount your kafka.keystore.jks and kafka.truststore.jks certs to the ${KAFKA_MOUNTED_CONF_DIR}/certs directory." + elif [[ "$KAFKA_TLS_TYPE" = "PEM" ]] && + { [[ ! -f "${KAFKA_CERTS_DIR}/kafka.keystore.pem" ]] || [[ ! -f "${KAFKA_CERTS_DIR}/kafka.keystore.key" ]] || [[ ! -f "$KAFKA_TLS_TRUSTSTORE_FILE" ]]; } && + { [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/certs/kafka.keystore.pem" ]] || [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/certs/kafka.keystore.key" ]] || [[ ! -f "$KAFKA_TLS_TRUSTSTORE_FILE" ]]; }; then + print_validation_error "In order to configure the TLS encryption for Kafka with PEM certs you must mount your kafka.keystore.pem, kafka.keystore.key and kafka.truststore.pem certs to the ${KAFKA_MOUNTED_CONF_DIR}/certs directory." + fi + fi + # If SASL/SASL_SSL listeners configured, check passwords are provided + if kafka_has_sasl_listener; then + if is_empty_value "${KAFKA_CFG_SASL_ENABLED_MECHANISMS:-}"; then + print_validation_error "Specified SASL protocol but no SASL mechanisms provided in KAFKA_CFG_SASL_ENABLED_MECHANISMS" + fi + fi + # Check users and passwords lists are the same size + read -r -a users <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_USERS:-}")" + read -r -a passwords <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_PASSWORDS:-}")" + if [[ "${#users[@]}" -ne "${#passwords[@]}" ]]; then + print_validation_error "Specify the same number of passwords on KAFKA_CLIENT_PASSWORDS as the number of users on KAFKA_CLIENT_USERS!" + fi + check_multi_value "KAFKA_TLS_TYPE" "JKS PEM" + check_multi_value "KAFKA_ZOOKEEPER_TLS_TYPE" "JKS PEM" + check_multi_value "KAFKA_ZOOKEEPER_PROTOCOL" "PLAINTEXT SASL SSL SASL_SSL" + check_multi_value "KAFKA_TLS_CLIENT_AUTH" "none requested required" + [[ "$error_code" -eq 0 ]] || return "$error_code" +} + +######################## +# Get kafka version +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# version +######################### +kafka_get_version() { + local -a cmd=("kafka-topics.sh" "--version") + am_i_root && cmd=("run_as_user" "$KAFKA_DAEMON_USER" "${cmd[@]}") + + read -r -a ver_split <<< "$("${cmd[@]}")" + echo "${ver_split[0]}" +} + +######################### +# Configure JAAS for a given listener and SASL mechanisms +# Globals: +# KAFKA_* +# Arguments: +# $1 - Name of the listener JAAS will be configured for +# $2 - Comma-separated list of SASL mechanisms to configure +# $3 - Comma-separated list of usernames +# $4 - Comma-separated list of passwords +# Returns: +# None +######################### +kafka_configure_server_jaas() { + local listener="${1:?missing listener name}" + local role="${2:-}" + + if [[ "$role" = "controller" ]]; then + local jaas_content=() + if [[ "$KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" = "PLAIN" ]]; then + jaas_content=( + "org.apache.kafka.common.security.plain.PlainLoginModule required" + "username=\"${KAFKA_CONTROLLER_USER}\"" + "password=\"${KAFKA_CONTROLLER_PASSWORD}\"" + "user_${KAFKA_CONTROLLER_USER}=\"${KAFKA_CONTROLLER_PASSWORD}\";" + ) + elif [[ "$KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" =~ SCRAM ]]; then + jaas_content=( + "org.apache.kafka.common.security.scram.ScramLoginModule required" + "username=\"${KAFKA_CONTROLLER_USER}\"" + "password=\"${KAFKA_CONTROLLER_PASSWORD}\";" + ) + fi + listener_lower="$(echo "$listener" | tr '[:upper:]' '[:lower:]')" + sasl_mechanism_lower="$(echo "$KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" | tr '[:upper:]' '[:lower:]')" + kafka_server_conf_set "listener.name.${listener_lower}.${sasl_mechanism_lower}.sasl.jaas.config" "${jaas_content[*]}" + else + read -r -a sasl_mechanisms_arr <<<"$(tr ',' ' ' <<<"$KAFKA_CFG_SASL_ENABLED_MECHANISMS")" + read -r -a users <<<"$(tr ',;' ' ' <<<"$KAFKA_CLIENT_USERS")" + read -r -a passwords <<<"$(tr ',;' ' ' <<<"$KAFKA_CLIENT_PASSWORDS")" + # Configure JAAS for each SASL mechanism + # ref: https://docs.confluent.io/platform/current/kafka/authentication_sasl/index.html + for sasl_mechanism in "${sasl_mechanisms_arr[@]}"; do + local jaas_content=() + # For PLAIN mechanism, only the first username will be used + if [[ "$sasl_mechanism" = "PLAIN" ]]; then + jaas_content=("org.apache.kafka.common.security.plain.PlainLoginModule required") + if [[ "$role" = "inter-broker" ]]; then + jaas_content+=( + "username=\"${KAFKA_INTER_BROKER_USER}\"" + "password=\"${KAFKA_INTER_BROKER_PASSWORD}\"" + ) + users+=("$KAFKA_INTER_BROKER_USER") + passwords+=("$KAFKA_INTER_BROKER_PASSWORD") + fi + for ((i = 0; i < ${#users[@]}; i++)); do + jaas_content+=("user_${users[i]}=\"${passwords[i]}\"") + done + # Add semi-colon to the last element of the array + jaas_content[${#jaas_content[@]} - 1]="${jaas_content[${#jaas_content[@]} - 1]};" + elif [[ "$sasl_mechanism" =~ SCRAM ]]; then + if [[ "$role" = "inter-broker" ]]; then + jaas_content=( + "org.apache.kafka.common.security.scram.ScramLoginModule required" + "username=\"${KAFKA_INTER_BROKER_USER}\"" + "password=\"${KAFKA_INTER_BROKER_PASSWORD}\";" + ) + else + jaas_content=("org.apache.kafka.common.security.scram.ScramLoginModule required;") + fi + fi + listener_lower="$(echo "$listener" | tr '[:upper:]' '[:lower:]')" + sasl_mechanism_lower="$(echo "$sasl_mechanism" | tr '[:upper:]' '[:lower:]')" + kafka_server_conf_set "listener.name.${listener_lower}.${sasl_mechanism_lower}.sasl.jaas.config" "${jaas_content[*]}" + done + fi +} + +######################## +# Configure Zookeeper JAAS authentication +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_zookeeper_configure_jaas(){ + local jaas_content=( + "org.apache.kafka.common.security.plain.PlainLoginModule required" + "username=\"${KAFKA_ZOOKEEPER_USER}\"" + "password=\"${KAFKA_ZOOKEEPER_PASSWORD}\";" + ) + + kafka_server_conf_set "sasl.jaas.config" "${jaas_content[*]}" +} + +######################## +# Generate JAAS authentication file for local producer/consumer to use +# Globals: +# KAFKA_* +# Arguments: +# $1 - Authentication protocol to use for the internal listener +# $2 - Authentication protocol to use for the client listener +# Returns: +# None +######################### +kafka_configure_consumer_producer_jaas(){ + local jaas_content=() + read -r -a users <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_USERS}")" + read -r -a passwords <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_PASSWORDS}")" + + if [[ "${KAFKA_CFG_SASL_ENABLED_MECHANISMS}" =~ SCRAM ]]; then + jaas_content=("org.apache.kafka.common.security.scram.ScramLoginModule required") + elif [[ "${KAFKA_CFG_SASL_ENABLED_MECHANISMS}" =~ PLAIN ]]; then + jaas_content=("org.apache.kafka.common.security.plain.PlainLoginModule required") + else + error "Couldn't configure a supported SASL mechanism for Kafka consumer/producer properties" + exit 1 + fi + + jaas_content+=( + "username=\"${users[0]}\"" + "password=\"${passwords[0]}\";" + ) + + kafka_producer_consumer_conf_set "sasl.jaas.config" "${jaas_content[*]}" +} + +######################## +# Create users in zookeper when using SASL/SCRAM mechanism +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_zookeeper_create_sasl_scram_users() { + info "Creating users in Zookeeper" + read -r -a users <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_USERS}")" + read -r -a passwords <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_PASSWORDS}")" + local zookeeper_connect + zookeeper_connect=$(grep "^zookeeper.connect=" "$KAFKA_CONF_FILE" | sed -E 's/^zookeeper\.connect=(\S+)$/\1/') + read -r -a zookeeper_hosts <<<"$(tr ',;' ' ' <<<"${zookeeper_connect}")" + + if [[ "${#zookeeper_hosts[@]}" -eq 0 ]]; then + error "Couldn't obtain zookeeper.connect from $KAFKA_CONF_FILE" + exit 1 + fi + # Wait for Zookeeper to be reachable + read -r -a aux <<<"$(tr ':' ' ' <<<"${zookeeper_hosts[0]}")" + local host="${aux[0]:?missing host}" + local port="${aux[1]:-2181}" + wait-for-port --host "$host" "$port" + + # Add interbroker credentials + if grep -Eq "^sasl.mechanism.inter.broker.protocol=SCRAM" "$KAFKA_CONF_FILE"; then + users+=("${KAFKA_INTER_BROKER_USER}") + passwords+=("${KAFKA_INTER_BROKER_PASSWORD}") + fi + for ((i = 0; i < ${#users[@]}; i++)); do + debug "Creating user ${users[i]} in zookeeper" + # Ref: https://docs.confluent.io/current/kafka/authentication_sasl/authentication_sasl_scram.html#sasl-scram-overview + debug_execute kafka-configs.sh --zookeeper "$zookeeper_connect" --alter --add-config "SCRAM-SHA-256=[iterations=8192,password=${passwords[i]}],SCRAM-SHA-512=[password=${passwords[i]}]" --entity-type users --entity-name "${users[i]}" + done +} + +######################## +# Configure Kafka SSL settings +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_configure_ssl() { + # Configures both Kafka server and producers/consumers + configure_both() { + kafka_server_conf_set "${1:?missing key}" "${2:?missing value}" + kafka_producer_consumer_conf_set "${1:?missing key}" "${2:?missing value}" + } + kafka_server_conf_set "ssl.client.auth" "${KAFKA_TLS_CLIENT_AUTH}" + configure_both ssl.keystore.type "${KAFKA_TLS_TYPE}" + configure_both ssl.truststore.type "${KAFKA_TLS_TYPE}" + local -r kafka_truststore_location="${KAFKA_CERTS_DIR}/$(basename "${KAFKA_TLS_TRUSTSTORE_FILE}")" + ! is_empty_value "${KAFKA_CERTIFICATE_PASSWORD:-}" && configure_both ssl.key.password "$KAFKA_CERTIFICATE_PASSWORD" + if [[ "$KAFKA_TLS_TYPE" = "PEM" ]]; then + file_to_multiline_property() { + awk 'NR > 1{print line"\\n\\"}{line=$0;}END{print $0" "}' <"${1:?missing file}" + } + remove_previous_cert_value() { + local key="${1:?missing key}" + files=( + "${KAFKA_CONF_FILE}" + "${KAFKA_CONF_DIR}/producer.properties" + "${KAFKA_CONF_DIR}/consumer.properties" + ) + for file in "${files[@]}"; do + if grep -q "^[#\\s]*$key\s*=.*" "$file"; then + # Delete all lines from the certificate beginning to its end + sed -i "/^[#\\s]*$key\s*=.*-----BEGIN/,/-----END/d" "$file" + fi + done + } + # We need to remove the previous cert value + # kafka_common_conf_set uses replace_in_file, which can't match multiple lines + remove_previous_cert_value ssl.keystore.key + remove_previous_cert_value ssl.keystore.certificate.chain + remove_previous_cert_value ssl.truststore.certificates + configure_both ssl.keystore.key "$(file_to_multiline_property "${KAFKA_CERTS_DIR}/kafka.keystore.key")" + configure_both ssl.keystore.certificate.chain "$(file_to_multiline_property "${KAFKA_CERTS_DIR}/kafka.keystore.pem")" + configure_both ssl.truststore.certificates "$(file_to_multiline_property "${kafka_truststore_location}")" + elif [[ "$KAFKA_TLS_TYPE" = "JKS" ]]; then + configure_both ssl.keystore.location "$KAFKA_CERTS_DIR"/kafka.keystore.jks + configure_both ssl.truststore.location "$kafka_truststore_location" + ! is_empty_value "${KAFKA_CERTIFICATE_PASSWORD:-}" && configure_both ssl.keystore.password "$KAFKA_CERTIFICATE_PASSWORD" + ! is_empty_value "${KAFKA_CERTIFICATE_PASSWORD:-}" && configure_both ssl.truststore.password "$KAFKA_CERTIFICATE_PASSWORD" + fi + true # Avoid the function to fail due to the check above +} + +######################## +# Get Zookeeper TLS settings +# Globals: +# KAFKA_ZOOKEEPER_TLS_* +# Arguments: +# None +# Returns: +# String +######################### +kafka_zookeeper_configure_tls() { + # Note that ZooKeeper does not support a key password different from the keystore password, + # so be sure to set the key password in the keystore to be identical to the keystore password; + # otherwise the connection attempt to Zookeeper will fail. + local keystore_location="" + local -r kafka_zk_truststore_location="${KAFKA_CERTS_DIR}/$(basename "${KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE}")" + + if [[ "$KAFKA_ZOOKEEPER_TLS_TYPE" = "JKS" ]] && [[ -f "${KAFKA_CERTS_DIR}/zookeeper.keystore.jks" ]]; then + keystore_location="${KAFKA_CERTS_DIR}/zookeeper.keystore.jks" + elif [[ "$KAFKA_ZOOKEEPER_TLS_TYPE" = "PEM" ]] && [[ -f "${KAFKA_CERTS_DIR}/zookeeper.keystore.pem" ]] && [[ -f "${KAFKA_CERTS_DIR}/zookeeper.keystore.key" ]]; then + # Concatenating private key into public certificate file + # This is needed to load keystore from location using PEM + keystore_location="${KAFKA_CERTS_DIR}/zookeeper.keypair.pem" + cat "${KAFKA_CERTS_DIR}/zookeeper.keystore.pem" "${KAFKA_CERTS_DIR}/zookeeper.keystore.key" > "$keystore_location" + fi + + kafka_server_conf_set "zookeeper.clientCnxnSocket" "org.apache.zookeeper.ClientCnxnSocketNetty" + kafka_server_conf_set "zookeeper.ssl.client.enable" "true" + is_boolean_yes "${KAFKA_ZOOKEEPER_TLS_VERIFY_HOSTNAME:-}" && kafka_server_conf_set "zookeeper.ssl.endpoint.identification.algorithm" "HTTPS" + ! is_empty_value "${keystore_location:-}" && kafka_server_conf_set "zookeeper.ssl.keystore.location" "${keystore_location}" + ! is_empty_value "${KAFKA_ZOOKEEPER_TLS_KEYSTORE_PASSWORD:-}" && kafka_server_conf_set "zookeeper.ssl.keystore.password" "${KAFKA_ZOOKEEPER_TLS_KEYSTORE_PASSWORD}" + ! is_empty_value "${kafka_zk_truststore_location:-}" && kafka_server_conf_set "zookeeper.ssl.truststore.location" "${kafka_zk_truststore_location}" + ! is_empty_value "${KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_PASSWORD:-}" && kafka_server_conf_set "zookeeper.ssl.truststore.password" "${KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_PASSWORD}" + true # Avoid the function to fail due to the check above +} + +######################## +# Configure Kafka configuration files from environment variables +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_configure_from_environment_variables() { + # List of special cases to apply to the variables + local -r exception_regexps=( + "s/sasl\.ssl/sasl_ssl/g" + "s/sasl\.plaintext/sasl_plaintext/g" + ) + # Map environment variables to config properties + for var in "${!KAFKA_CFG_@}"; do + key="$(echo "$var" | sed -e 's/^KAFKA_CFG_//g' -e 's/_/\./g' | tr '[:upper:]' '[:lower:]')" + + # Exception for the camel case in this environment variable + [[ "$var" == "KAFKA_CFG_ZOOKEEPER_CLIENTCNXNSOCKET" ]] && key="zookeeper.clientCnxnSocket" + + # Apply exception regexps + for regex in "${exception_regexps[@]}"; do + key="$(echo "$key" | sed "$regex")" + done + + value="${!var}" + kafka_server_conf_set "$key" "$value" + done +} + +######################## +# Initialize KRaft storage +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_kraft_storage_initialize() { + local args=("--config" "$KAFKA_CONF_FILE" "--ignore-formatted") + info "Initializing KRaft storage metadata" + + # If cluster.id found in meta.properties, use it + if [[ -f "${KAFKA_DATA_DIR}/meta.properties" ]]; then + KAFKA_KRAFT_CLUSTER_ID=$(grep "^cluster.id=" "${KAFKA_DATA_DIR}/meta.properties" | sed -E 's/^cluster\.id=(\S+)$/\1/') + fi + + if is_empty_value "${KAFKA_KRAFT_CLUSTER_ID:-}"; then + warn "KAFKA_KRAFT_CLUSTER_ID not set - If using multiple nodes then you must use the same Cluster ID for each one" + KAFKA_KRAFT_CLUSTER_ID="$("${KAFKA_HOME}/bin/kafka-storage.sh" random-uuid)" + info "Generated Kafka cluster ID '${KAFKA_KRAFT_CLUSTER_ID}'" + fi + args+=("--cluster-id=$KAFKA_KRAFT_CLUSTER_ID") + + # SCRAM users are configured during the cluster bootstrapping process and can later be manually updated using kafka-config.sh + if is_boolean_yes "${KAFKA_KRAFT_BOOTSTRAP_SCRAM_USERS:-}"; then + info "Adding KRaft SCRAM users at storage bootstrap" + read -r -a users <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_USERS}")" + read -r -a passwords <<<"$(tr ',;' ' ' <<<"${KAFKA_CLIENT_PASSWORDS}")" + # Configure SCRAM-SHA-256 if enabled + if grep -Eq "^sasl.enabled.mechanisms=.*SCRAM-SHA-256" "$KAFKA_CONF_FILE"; then + for ((i = 0; i < ${#users[@]}; i++)); do + args+=("--add-scram" "SCRAM-SHA-256=[name=${users[i]},password=${passwords[i]}]") + done + fi + # Configure SCRAM-SHA-512 if enabled + if grep -Eq "^sasl.enabled.mechanisms=.*SCRAM-SHA-512" "$KAFKA_CONF_FILE"; then + for ((i = 0; i < ${#users[@]}; i++)); do + args+=("--add-scram" "SCRAM-SHA-512=[name=${users[i]},password=${passwords[i]}]") + done + fi + # Add interbroker credentials + if grep -Eq "^sasl.mechanism.inter.broker.protocol=SCRAM-SHA-256" "$KAFKA_CONF_FILE"; then + args+=("--add-scram" "SCRAM-SHA-256=[name=${KAFKA_INTER_BROKER_USER},password=${KAFKA_INTER_BROKER_PASSWORD}]") + elif grep -Eq "^sasl.mechanism.inter.broker.protocol=SCRAM-SHA-512" "$KAFKA_CONF_FILE"; then + args+=("--add-scram" "SCRAM-SHA-512=[name=${KAFKA_INTER_BROKER_USER},password=${KAFKA_INTER_BROKER_PASSWORD}]") + fi + # Add controller credentials + if grep -Eq "^sasl.mechanism.controller.protocol=SCRAM-SHA-256" "$KAFKA_CONF_FILE"; then + args+=("--add-scram" "SCRAM-SHA-256=[name=${KAFKA_CONTROLLER_USER},password=${KAFKA_CONTROLLER_PASSWORD}]") + elif grep -Eq "^sasl.mechanism.controller.protocol=SCRAM-SHA-512" "$KAFKA_CONF_FILE"; then + args+=("--add-scram" "SCRAM-SHA-512=[name=${KAFKA_CONTROLLER_USER},password=${KAFKA_CONTROLLER_PASSWORD}]") + fi + fi + info "Formatting storage directories to add metadata..." + "${KAFKA_HOME}/bin/kafka-storage.sh" format "${args[@]}" +} + +######################## +# Detects inconsitences between the configuration at KAFKA_CONF_FILE and cluster-state file +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_kraft_quorum_voters_changed(){ + read -r -a quorum_voters_conf_ids <<<"$(grep "^controller.quorum.voters=" "$KAFKA_CONF_FILE" | sed "s/^controller.quorum.voters=//" | tr "," " " | sed -E "s/\@\S+//g")" + read -r -a quorum_voters_state_ids <<< "$(grep -Eo "\{\"voterId\":[0-9]+\}" "${KAFKA_DATA_DIR}/__cluster_metadata-0/quorum-state" | grep -Eo "[0-9]+" | tr "\n" " ")" + + if [[ "${#quorum_voters_conf_ids[@]}" != "${#quorum_voters_state_ids[@]}" ]]; then + true + else + read -r -a sorted_state <<< "$(echo "${quorum_voters_conf_ids[@]}" | tr ' ' '\n' | sort | tr '\n' ' ')" + read -r -a sorted_conf <<< "$(echo "${quorum_voters_state_ids[@]}" | tr ' ' '\n' | sort | tr '\n' ' ')" + if [[ "${sorted_state[*]}" = "${sorted_conf[*]}" ]]; then + false + else + true + fi + fi +} + +######################## +# Initialize Kafka +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_initialize() { + info "Initializing Kafka..." + # Check for mounted configuration files + if ! is_dir_empty "$KAFKA_MOUNTED_CONF_DIR"; then + cp -Lr "$KAFKA_MOUNTED_CONF_DIR"/* "$KAFKA_CONF_DIR" + fi + # Copy truststore to cert directory + for cert_var in KAFKA_TLS_TRUSTSTORE_FILE KAFKA_ZOOKEEPER_TLS_TRUSTSTORE_FILE; do + # Only copy if the file exists and it is in a different location than KAFKA_CERTS_DIR (to avoid copying to the same location) + if [[ -f "${!cert_var}" ]] && ! [[ "${!cert_var}" =~ $KAFKA_CERTS_DIR ]]; then + info "Copying truststore ${!cert_var} to ${KAFKA_CERTS_DIR}" + cp -L "${!cert_var}" "$KAFKA_CERTS_DIR" + fi + done + + if [[ ! -f "${KAFKA_MOUNTED_CONF_DIR}/server.properties" ]]; then + info "No injected configuration files found, creating default config files" + # Restore original server.properties but remove Zookeeper/KRaft specific settings for compatibility with both architectures + cp "${KAFKA_CONF_DIR}/server.properties.original" "$KAFKA_CONF_FILE" + kafka_server_unify_conf + # Configure Kafka settings + kafka_server_conf_set log.dirs "$KAFKA_DATA_DIR" + kafka_configure_from_environment_variables + # Configure Kafka producer/consumer to set up message sizes + ! is_empty_value "${KAFKA_CFG_MAX_REQUEST_SIZE:-}" && kafka_common_conf_set "$KAFKA_CONF_DIR/producer.properties" max.request.size "$KAFKA_CFG_MAX_REQUEST_SIZE" + ! is_empty_value "${KAFKA_CFG_MAX_PARTITION_FETCH_BYTES:-}" && kafka_common_conf_set "$KAFKA_CONF_DIR/consumer.properties" max.partition.fetch.bytes "$KAFKA_CFG_MAX_PARTITION_FETCH_BYTES" + # Zookeeper mode additional settings + if ! is_empty_value "${KAFKA_CFG_ZOOKEEPER_CONNECT:-}"; then + if [[ "$KAFKA_ZOOKEEPER_PROTOCOL" =~ SSL ]]; then + kafka_zookeeper_configure_tls + fi + if [[ "$KAFKA_ZOOKEEPER_PROTOCOL" =~ SASL ]]; then + kafka_zookeeper_configure_jaas + fi + fi + # If at least one listener uses SSL or SASL_SSL, ensure SSL is configured + if kafka_has_ssl_listener; then + kafka_configure_ssl + fi + # If at least one listener uses SASL_PLAINTEXT or SASL_SSL, ensure SASL is configured + if kafka_has_sasl_listener; then + if [[ "$KAFKA_CFG_SASL_ENABLED_MECHANISMS" =~ SCRAM ]]; then + if ! is_empty_value "${KAFKA_CFG_PROCESS_ROLES:-}"; then + if [[ "$(kafka_get_version)" =~ ^3\.2\.|^3\.3\.|^3\.4\. ]]; then + # NOTE: This will depend on Kafka version when support for SCRAM is added + warn "KRaft mode requires Kafka version 3.5 or higher for SCRAM to be supported. SCRAM SASL mechanisms will now be disabled." + KAFKA_CFG_SASL_ENABLED_MECHANISMS=PLAIN + else + export KAFKA_KRAFT_BOOTSTRAP_SCRAM_USERS="true" + fi + fi + if ! is_empty_value "${KAFKA_CFG_ZOOKEEPER_CONNECT:-}"; then + export KAFKA_ZOOKEEPER_BOOTSTRAP_SCRAM_USERS="true" + fi + fi + kafka_server_conf_set sasl.enabled.mechanisms "$KAFKA_CFG_SASL_ENABLED_MECHANISMS" + fi + # Settings for each Kafka Listener are configured individually + read -r -a protocol_maps <<<"$(tr ',' ' ' <<<"$KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP")" + for protocol_map in "${protocol_maps[@]}"; do + read -r -a map <<<"$(tr ':' ' ' <<<"$protocol_map")" + # Obtain the listener and protocol from protocol map string, e.g. CONTROLLER:PLAINTEXT + listener="${map[0]}" + protocol="${map[1]}" + listener_lower="$(echo "$listener" | tr '[:upper:]' '[:lower:]')" + + if [[ "$protocol" = "SSL" || "$protocol" = "SASL_SSL" ]]; then + listener_upper="$(echo "$listener" | tr '[:lower:]' '[:upper:]')" + env_name="KAFKA_TLS_${listener_upper}_CLIENT_AUTH" + [[ -n "${!env_name:-}" ]] && kafka_server_conf_set "listener.name.${listener_lower}.ssl.client.auth" "${!env_name}" + fi + if [[ "$protocol" = "SASL_PLAINTEXT" || "$protocol" = "SASL_SSL" ]]; then + local role="" + if [[ "$listener" = "${KAFKA_CFG_INTER_BROKER_LISTENER_NAME:-INTERNAL}" ]]; then + kafka_server_conf_set sasl.mechanism.inter.broker.protocol "$KAFKA_CFG_SASL_MECHANISM_INTER_BROKER_PROTOCOL" + role="inter-broker" + elif [[ "${KAFKA_CFG_CONTROLLER_LISTENER_NAMES:-CONTROLLER}" =~ $listener ]]; then + kafka_server_conf_set sasl.mechanism.controller.protocol "$KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" + kafka_server_conf_set "listener.name.${listener_lower}.sasl.enabled.mechanisms" "$KAFKA_CFG_SASL_MECHANISM_CONTROLLER_PROTOCOL" + role="controller" + fi + # If KAFKA_CLIENT_LISTENER_NAME is found in the listeners list, configure the producer/consumer accordingly + if [[ "$listener" = "${KAFKA_CLIENT_LISTENER_NAME:-CLIENT}" ]]; then + kafka_configure_consumer_producer_jaas + kafka_producer_consumer_conf_set security.protocol "$protocol" + kafka_producer_consumer_conf_set sasl.mechanism "${KAFKA_CLIENT_SASL_MECHANISM:-$(kafka_client_sasl_mechanism)}" + fi + # Configure inline listener jaas configuration, omitted if mounted JAAS conf file detected + if [[ ! -f "${KAFKA_CONF_DIR}/kafka_jaas.conf" ]]; then + kafka_configure_server_jaas "$listener_lower" "${role:-}" + fi + fi + done + # Configure Kafka using environment variables + # This is executed at the end, to allow users to override properties set by the initialization logic + kafka_configure_from_environment_variables + else + info "Detected mounted server.properties file at ${KAFKA_MOUNTED_CONF_DIR}/server.properties. Skipping configuration based on env variables" + fi + true +} + +######################## +# Returns the most secure SASL mechanism available for Kafka clients +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################## +kafka_client_sasl_mechanism() { + local sasl_mechanism="" + + if [[ "$KAFKA_CFG_SASL_ENABLED_MECHANISMS" =~ SCRAM-SHA-512 ]]; then + sasl_mechanism="SCRAM-SHA-512" + elif [[ "$KAFKA_CFG_SASL_ENABLED_MECHANISMS" =~ SCRAM-SHA-256 ]]; then + sasl_mechanism="SCRAM-SHA-256" + elif [[ "$KAFKA_CFG_SASL_ENABLED_MECHANISMS" =~ PLAIN ]]; then + sasl_mechanism="PLAIN" + fi + echo "$sasl_mechanism" +} + +######################## +# Removes default settings referencing Zookeeper mode or KRaft mode +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################## +kafka_server_unify_conf() { + local -r remove_regexps=( + #Zookeeper + "s/^zookeeper\./#zookeeper./g" + "s/^group\.initial/#group.initial/g" + "s/^broker\./#broker./g" + "s/^node\./#node./g" + "s/^process\./#process./g" + "s/^listeners=/#listeners=/g" + "s/^listener\./#listener./g" + "s/^controller\./#controller./g" + "s/^inter\.broker/#inter.broker/g" + "s/^advertised\.listeners/#advertised.listeners/g" + ) + + # Map environment variables to config properties + for regex in "${remove_regexps[@]}"; do + sed -i "${regex}" "$KAFKA_CONF_FILE" + done +} + +######################## +# Dinamically set node.id/broker.id/controller.quorum.voters if their alternative environment variable _COMMAND is set +# Globals: +# KAFKA_*_COMMAND +# Arguments: +# None +# Returns: +# None +######################### +kafka_dynamic_environment_variables() { + # KRaft mode + if ! is_empty_value "${KAFKA_NODE_ID_COMMAND:-}"; then + KAFKA_CFG_NODE_ID="$(eval "${KAFKA_NODE_ID_COMMAND}")" + export KAFKA_CFG_NODE_ID + fi + if ! is_empty_value "${KAFKA_CONTROLLER_QUORUM_VOTERS_COMMAND:-}"; then + KAFKA_CFG_CONTROLLER_QUORUM_VOTERS="$(eval "${KAFKA_CONTROLLER_QUORUM_VOTERS_COMMAND}")" + export KAFKA_CFG_CONTROLLER_QUORUM_VOTERS + fi + # Zookeeper mode + # DEPRECATED - BROKER_ID_COMMAND has been deprecated, please use KAFKA_BROKER_ID_COMMAND instead + if ! is_empty_value "${KAFKA_BROKER_ID_COMMAND:-}"; then + KAFKA_CFG_BROKER_ID="$(eval "${KAFKA_BROKER_ID_COMMAND}")" + export KAFKA_CFG_BROKER_ID + elif ! is_empty_value "${BROKER_ID_COMMAND:-}"; then + KAFKA_CFG_BROKER_ID="$(eval "${BROKER_ID_COMMAND}")" + export KAFKA_CFG_BROKER_ID + fi +} + +######################## +# Run custom initialization scripts +# Globals: +# KAFKA_* +# Arguments: +# None +# Returns: +# None +######################### +kafka_custom_init_scripts() { + if [[ -n $(find "${KAFKA_INITSCRIPTS_DIR}/" -type f -regex ".*\.\(sh\)") ]] && [[ ! -f "${KAFKA_VOLUME_DIR}/.user_scripts_initialized" ]]; then + info "Loading user's custom files from $KAFKA_INITSCRIPTS_DIR" + for f in /docker-entrypoint-initdb.d/*; do + debug "Executing $f" + case "$f" in + *.sh) + if [[ -x "$f" ]]; then + if ! "$f"; then + error "Failed executing $f" + return 1 + fi + else + warn "Sourcing $f as it is not executable by the current user, any error may cause initialization to fail" + . "$f" + fi + ;; + *) + warn "Skipping $f, supported formats are: .sh" + ;; + esac + done + touch "$KAFKA_VOLUME_DIR"/.user_scripts_initialized + fi +} + +######################## +# Check if Kafka is running +# Globals: +# KAFKA_PID_FILE +# Arguments: +# None +# Returns: +# Whether Kafka is running +######################## +is_kafka_running() { + local pid + pid="$(get_pid_from_file "$KAFKA_PID_FILE")" + if [[ -n "$pid" ]]; then + is_service_running "$pid" + else + false + fi +} + +######################## +# Check if Kafka is running +# Globals: +# KAFKA_PID_FILE +# Arguments: +# None +# Returns: +# Whether Kafka is not running +######################## +is_kafka_not_running() { + ! is_kafka_running +} + +######################## +# Stop Kafka +# Globals: +# KAFKA_PID_FILE +# Arguments: +# None +# Returns: +# None +######################### +kafka_stop() { + ! is_kafka_running && return + stop_service_using_pid "$KAFKA_PID_FILE" TERM +} diff --git a/container/bitnami/tags-info.yaml b/container/bitnami/tags-info.yaml new file mode 100644 index 0000000000..7f7db72ac1 --- /dev/null +++ b/container/bitnami/tags-info.yaml @@ -0,0 +1,5 @@ +rolling-tags: +- "3.9" +- 3.9-debian-12 +- 3.9.0 +- latest diff --git a/core/src/main/java/kafka/autobalancer/AutoBalancerListener.java b/core/src/main/java/kafka/autobalancer/AutoBalancerListener.java index 75ff94ca98..2b68ccad8e 100644 --- a/core/src/main/java/kafka/autobalancer/AutoBalancerListener.java +++ b/core/src/main/java/kafka/autobalancer/AutoBalancerListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer; diff --git a/core/src/main/java/kafka/autobalancer/AutoBalancerManager.java b/core/src/main/java/kafka/autobalancer/AutoBalancerManager.java index 3127a80387..7be45102bf 100644 --- a/core/src/main/java/kafka/autobalancer/AutoBalancerManager.java +++ b/core/src/main/java/kafka/autobalancer/AutoBalancerManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer; @@ -75,6 +83,7 @@ protected void init() { this.anomalyDetector = new AnomalyDetectorImpl(config.originals(), new LogContext(String.format("[AnomalyDetector id=%d] ", nodeId)), clusterModel, actionExecutorService); + ((AnomalyDetectorImpl) this.anomalyDetector).lockedNodes(() -> quorumController.nodeControlManager().lockedNodes()); this.reconfigurables.add(anomalyDetector); diff --git a/core/src/main/java/kafka/autobalancer/LoadRetriever.java b/core/src/main/java/kafka/autobalancer/LoadRetriever.java index 7db48ad765..666b2cbdcc 100644 --- a/core/src/main/java/kafka/autobalancer/LoadRetriever.java +++ b/core/src/main/java/kafka/autobalancer/LoadRetriever.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer; @@ -49,8 +57,10 @@ import org.apache.kafka.common.serialization.StringDeserializer; import org.apache.kafka.controller.Controller; import org.apache.kafka.controller.ControllerRequestContext; +import org.apache.kafka.server.config.QuotaConfigs; import com.automq.stream.utils.LogContext; +import com.automq.stream.utils.Threads; import java.time.Duration; import java.util.Collections; @@ -64,7 +74,6 @@ import java.util.Random; import java.util.Set; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; @@ -84,7 +93,7 @@ public class LoadRetriever extends AbstractResumableService implements BrokerSta private final Condition cond; private final Controller controller; private final ScheduledExecutorService mainExecutorService; - private final Set brokerIdsInUse; + private final Map bootstrapServerMapInUse; private final Set currentAssignment = new HashSet<>(); private final StaticAutoBalancerConfig staticConfig; private final String listenerName; @@ -101,17 +110,19 @@ public LoadRetriever(AutoBalancerControllerConfig config, Controller controller, this.controller = controller; this.clusterModel = clusterModel; this.bootstrapServerMap = new HashMap<>(); - this.brokerIdsInUse = new HashSet<>(); + this.bootstrapServerMapInUse = new HashMap<>(); this.lock = new ReentrantLock(); this.cond = lock.newCondition(); - this.mainExecutorService = Executors.newSingleThreadScheduledExecutor(new AutoBalancerThreadFactory("load-retriever-main")); + this.mainExecutorService = + Threads.newSingleThreadScheduledExecutor( + new AutoBalancerThreadFactory("load-retriever-main"), logger); leaderEpochInitialized = false; staticConfig = new StaticAutoBalancerConfig(config.originals(), false); listenerName = staticConfig.getString(StaticAutoBalancerConfig.AUTO_BALANCER_CLIENT_LISTENER_NAME_CONFIG); metricReporterTopicPartition = config.getInt(AutoBalancerControllerConfig.AUTO_BALANCER_CONTROLLER_METRICS_TOPIC_NUM_PARTITIONS_CONFIG); metricReporterTopicRetentionTime = config.getLong(AutoBalancerControllerConfig.AUTO_BALANCER_CONTROLLER_METRICS_TOPIC_RETENTION_MS_CONFIG); consumerPollTimeout = config.getLong(AutoBalancerControllerConfig.AUTO_BALANCER_CONTROLLER_CONSUMER_POLL_TIMEOUT); - consumerClientIdPrefix = config.getString(AutoBalancerControllerConfig.AUTO_BALANCER_CONTROLLER_CONSUMER_CLIENT_ID_PREFIX); + consumerClientIdPrefix = QuotaConfigs.INTERNAL_CLIENT_ID_PREFIX + config.getString(AutoBalancerControllerConfig.AUTO_BALANCER_CONTROLLER_CONSUMER_CLIENT_ID_PREFIX); consumerRetryBackOffMs = config.getLong(AutoBalancerControllerConfig.AUTO_BALANCER_CONTROLLER_CONSUMER_RETRY_BACKOFF_MS); } @@ -154,7 +165,7 @@ protected Properties buildConsumerProps(String bootstrapServer) { Properties consumerProps = new Properties(); long randomToken = RANDOM.nextLong(); consumerProps.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer); - consumerProps.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, consumerClientIdPrefix + "-consumer-" + randomToken); + consumerProps.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, consumerClientIdPrefix + randomToken); consumerProps.setProperty(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, Long.toString(consumerRetryBackOffMs)); consumerProps.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest"); consumerProps.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false"); @@ -167,8 +178,8 @@ protected Properties buildConsumerProps(String bootstrapServer) { public static class BrokerEndpoints { private final int brokerId; private Set endpoints = new HashSet<>(); - private boolean isFenced; + private boolean isOutdated = false; public BrokerEndpoints(int brokerId) { this.brokerId = brokerId; @@ -199,6 +210,13 @@ public boolean isValid() { return !this.isFenced && !this.endpoints.isEmpty(); } + public boolean isOutdated() { + return isOutdated; + } + + public void setOutdated(boolean outdated) { + isOutdated = outdated; + } } @Override @@ -221,7 +239,13 @@ public void onBrokerRegister(RegisterBrokerRecord record) { BrokerEndpoints brokerEndpoints = new BrokerEndpoints(record.brokerId()); brokerEndpoints.setFenced(Utils.isBrokerFenced(record)); brokerEndpoints.setEndpoints(endpoints); + brokerEndpoints.setOutdated(false); this.bootstrapServerMap.put(record.brokerId(), brokerEndpoints); + this.bootstrapServerMapInUse.computeIfPresent(record.brokerId(), (k, v) -> { + v.setOutdated(!v.getEndpoints().equals(endpoints)); + v.setFenced(Utils.isBrokerFenced(record)); + return v; + }); cond.signal(); } finally { lock.unlock(); @@ -236,7 +260,6 @@ public void onBrokerUnregister(UnregisterBrokerRecord record) { } finally { lock.unlock(); } - } @Override @@ -245,10 +268,14 @@ public void onBrokerRegistrationChanged(BrokerRegistrationChangeRecord record) { isBrokerFenced.ifPresent(isFenced -> { lock.lock(); try { - BrokerEndpoints brokerEndpoints = this.bootstrapServerMap.get(record.brokerId()); - if (brokerEndpoints != null) { - brokerEndpoints.setFenced(isFenced); - } + this.bootstrapServerMap.computeIfPresent(record.brokerId(), (k, v) -> { + v.setFenced(isFenced); + return v; + }); + this.bootstrapServerMapInUse.computeIfPresent(record.brokerId(), (k, v) -> { + v.setFenced(isFenced); + return v; + }); cond.signal(); } finally { lock.unlock(); @@ -256,20 +283,21 @@ public void onBrokerRegistrationChanged(BrokerRegistrationChangeRecord record) { }); } - private boolean hasAvailableBrokerInUse() { - if (brokerIdsInUse.isEmpty()) { + boolean hasAvailableBrokerInUse() { + if (bootstrapServerMapInUse.isEmpty()) { return false; } - for (int brokerId : brokerIdsInUse) { - BrokerEndpoints brokerEndpoints = this.bootstrapServerMap.get(brokerId); - if (brokerEndpoints != null && brokerEndpoints.isValid()) { + for (Map.Entry entry : bootstrapServerMapInUse.entrySet()) { + int brokerId = entry.getKey(); + BrokerEndpoints endpoints = entry.getValue(); + if (bootstrapServerMap.containsKey(brokerId) && endpoints != null && endpoints.isValid() && !endpoints.isOutdated()) { return true; } } return false; } - private boolean hasAvailableBroker() { + boolean hasAvailableBroker() { if (this.bootstrapServerMap.isEmpty()) { return false; } @@ -283,18 +311,17 @@ private boolean hasAvailableBroker() { public String buildBootstrapServer() { Set endpoints = new HashSet<>(); - this.brokerIdsInUse.clear(); + this.bootstrapServerMapInUse.clear(); for (BrokerEndpoints brokerEndpoints : this.bootstrapServerMap.values()) { if (brokerEndpoints.isValid() && !brokerEndpoints.getEndpoints().isEmpty()) { endpoints.add(brokerEndpoints.getEndpoints().iterator().next()); - this.brokerIdsInUse.add(brokerEndpoints.brokerId()); + this.bootstrapServerMapInUse.put(brokerEndpoints.brokerId(), brokerEndpoints); } } return String.join(",", endpoints); } - private void checkAndCreateConsumer(int epoch) { - String bootstrapServer; + void checkAndCreateConsumer(int epoch) { this.lock.lock(); try { if (!isRunnable(epoch)) { @@ -314,9 +341,10 @@ private void checkAndCreateConsumer(int epoch) { return; } } - bootstrapServer = buildBootstrapServer(); - if (this.consumer == null && !bootstrapServer.isEmpty()) { + + if (this.consumer == null) { //TODO: fetch metadata from controller + String bootstrapServer = buildBootstrapServer(); this.consumer = createConsumer(bootstrapServer); logger.info("Created consumer on {}", bootstrapServer); } diff --git a/core/src/main/java/kafka/autobalancer/common/Action.java b/core/src/main/java/kafka/autobalancer/common/Action.java index 88899eba45..8333c4905e 100644 --- a/core/src/main/java/kafka/autobalancer/common/Action.java +++ b/core/src/main/java/kafka/autobalancer/common/Action.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common; diff --git a/core/src/main/java/kafka/autobalancer/common/ActionType.java b/core/src/main/java/kafka/autobalancer/common/ActionType.java index 72949f1faf..a8ec5f5881 100644 --- a/core/src/main/java/kafka/autobalancer/common/ActionType.java +++ b/core/src/main/java/kafka/autobalancer/common/ActionType.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common; diff --git a/core/src/main/java/kafka/autobalancer/common/AutoBalancerConstants.java b/core/src/main/java/kafka/autobalancer/common/AutoBalancerConstants.java index e1bfd79104..92c047fe47 100644 --- a/core/src/main/java/kafka/autobalancer/common/AutoBalancerConstants.java +++ b/core/src/main/java/kafka/autobalancer/common/AutoBalancerConstants.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common; diff --git a/core/src/main/java/kafka/autobalancer/common/AutoBalancerThreadFactory.java b/core/src/main/java/kafka/autobalancer/common/AutoBalancerThreadFactory.java index 7c9570f7c9..c7b2a73800 100644 --- a/core/src/main/java/kafka/autobalancer/common/AutoBalancerThreadFactory.java +++ b/core/src/main/java/kafka/autobalancer/common/AutoBalancerThreadFactory.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common; diff --git a/core/src/main/java/kafka/autobalancer/common/Utils.java b/core/src/main/java/kafka/autobalancer/common/Utils.java index 448044bc29..23eb073ee9 100644 --- a/core/src/main/java/kafka/autobalancer/common/Utils.java +++ b/core/src/main/java/kafka/autobalancer/common/Utils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common; diff --git a/core/src/main/java/kafka/autobalancer/common/normalizer/LinearNormalizer.java b/core/src/main/java/kafka/autobalancer/common/normalizer/LinearNormalizer.java index 5c14559158..31a781e331 100644 --- a/core/src/main/java/kafka/autobalancer/common/normalizer/LinearNormalizer.java +++ b/core/src/main/java/kafka/autobalancer/common/normalizer/LinearNormalizer.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.normalizer; diff --git a/core/src/main/java/kafka/autobalancer/common/normalizer/Normalizer.java b/core/src/main/java/kafka/autobalancer/common/normalizer/Normalizer.java index e579f1b9e4..48e3734517 100644 --- a/core/src/main/java/kafka/autobalancer/common/normalizer/Normalizer.java +++ b/core/src/main/java/kafka/autobalancer/common/normalizer/Normalizer.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.normalizer; diff --git a/core/src/main/java/kafka/autobalancer/common/normalizer/StepNormalizer.java b/core/src/main/java/kafka/autobalancer/common/normalizer/StepNormalizer.java index f10f329523..f45840aae3 100644 --- a/core/src/main/java/kafka/autobalancer/common/normalizer/StepNormalizer.java +++ b/core/src/main/java/kafka/autobalancer/common/normalizer/StepNormalizer.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.normalizer; diff --git a/core/src/main/java/kafka/autobalancer/common/types/MetricTypes.java b/core/src/main/java/kafka/autobalancer/common/types/MetricTypes.java index ea7677e8d1..e5c82bf79a 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/MetricTypes.java +++ b/core/src/main/java/kafka/autobalancer/common/types/MetricTypes.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types; diff --git a/core/src/main/java/kafka/autobalancer/common/types/MetricVersion.java b/core/src/main/java/kafka/autobalancer/common/types/MetricVersion.java index 7b3fc1a3b5..d5e198eed5 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/MetricVersion.java +++ b/core/src/main/java/kafka/autobalancer/common/types/MetricVersion.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types; diff --git a/core/src/main/java/kafka/autobalancer/common/types/RawMetricTypes.java b/core/src/main/java/kafka/autobalancer/common/types/RawMetricTypes.java index d88c746cc0..a9d21293a3 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/RawMetricTypes.java +++ b/core/src/main/java/kafka/autobalancer/common/types/RawMetricTypes.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types; @@ -25,7 +33,7 @@ public class RawMetricTypes { public static final byte BROKER_MAX_PENDING_FETCH_LATENCY_MS = (byte) 5; public static final byte BROKER_METRIC_VERSION = (byte) 6; public static final Map ABNORMAL_METRICS = Map.of( - BROKER_APPEND_LATENCY_AVG_MS, new AbnormalLatency(100), // 100ms +// BROKER_APPEND_LATENCY_AVG_MS, new AbnormalLatency(100), // 100ms BROKER_MAX_PENDING_APPEND_LATENCY_MS, new AbnormalLatency(10000), // 10s BROKER_MAX_PENDING_FETCH_LATENCY_MS, new AbnormalLatency(10000) // 10s ); diff --git a/core/src/main/java/kafka/autobalancer/common/types/Resource.java b/core/src/main/java/kafka/autobalancer/common/types/Resource.java index 449955f699..08bc7c4bb9 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/Resource.java +++ b/core/src/main/java/kafka/autobalancer/common/types/Resource.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types; diff --git a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalLatency.java b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalLatency.java index 5480450a4f..45f7e50276 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalLatency.java +++ b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalLatency.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types.metrics; diff --git a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalMetric.java b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalMetric.java index 7a7c40ee71..f1c06fce5c 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalMetric.java +++ b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalMetric.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types.metrics; diff --git a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalQueueSize.java b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalQueueSize.java index 9375680177..8b8b8cee0e 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalQueueSize.java +++ b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbnormalQueueSize.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types.metrics; diff --git a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbstractSimpleAbnormalMetric.java b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbstractSimpleAbnormalMetric.java index e7a4b6566a..9ddd2d046b 100644 --- a/core/src/main/java/kafka/autobalancer/common/types/metrics/AbstractSimpleAbnormalMetric.java +++ b/core/src/main/java/kafka/autobalancer/common/types/metrics/AbstractSimpleAbnormalMetric.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.types.metrics; diff --git a/core/src/main/java/kafka/autobalancer/config/AutoBalancerControllerConfig.java b/core/src/main/java/kafka/autobalancer/config/AutoBalancerControllerConfig.java index 96057c141e..ed606730dc 100644 --- a/core/src/main/java/kafka/autobalancer/config/AutoBalancerControllerConfig.java +++ b/core/src/main/java/kafka/autobalancer/config/AutoBalancerControllerConfig.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.config; @@ -53,7 +61,7 @@ public class AutoBalancerControllerConfig extends AbstractConfig { public static final Integer DEFAULT_AUTO_BALANCER_CONTROLLER_METRICS_TOPIC_NUM_PARTITIONS_CONFIG = 1; public static final long DEFAULT_AUTO_BALANCER_CONTROLLER_METRICS_TOPIC_RETENTION_MS_CONFIG = TimeUnit.MINUTES.toMillis(30); public static final long DEFAULT_AUTO_BALANCER_CONTROLLER_CONSUMER_POLL_TIMEOUT = 1000L; - public static final String DEFAULT_AUTO_BALANCER_CONTROLLER_CONSUMER_CLIENT_ID_PREFIX = "AutoBalancerControllerConsumer"; + public static final String DEFAULT_AUTO_BALANCER_CONTROLLER_CONSUMER_CLIENT_ID_PREFIX = "auto_balancer_controller_consumer_"; public static final long DEFAULT_AUTO_BALANCER_CONTROLLER_CONSUMER_RETRY_BACKOFF_MS = 1000; public static final long DEFAULT_AUTO_BALANCER_CONTROLLER_ACCEPTED_METRICS_DELAY_MS = Duration.ofMinutes(1).toMillis(); public static final String DEFAULT_AUTO_BALANCER_CONTROLLER_GOALS = new StringJoiner(",") diff --git a/core/src/main/java/kafka/autobalancer/config/AutoBalancerMetricsReporterConfig.java b/core/src/main/java/kafka/autobalancer/config/AutoBalancerMetricsReporterConfig.java index 4819274dd8..ef22c0213d 100644 --- a/core/src/main/java/kafka/autobalancer/config/AutoBalancerMetricsReporterConfig.java +++ b/core/src/main/java/kafka/autobalancer/config/AutoBalancerMetricsReporterConfig.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.config; @@ -36,7 +44,7 @@ public class AutoBalancerMetricsReporterConfig extends AbstractConfig { public static final String AUTO_BALANCER_METRICS_REPORTER_LINGER_MS_CONFIG = PREFIX + ProducerConfig.LINGER_MS_CONFIG; public static final String AUTO_BALANCER_METRICS_REPORTER_BATCH_SIZE_CONFIG = PREFIX + ProducerConfig.BATCH_SIZE_CONFIG; /* Default values */ - public static final String DEFAULT_AUTO_BALANCER_METRICS_REPORTER_PRODUCER_CLIENT_ID = "AutoBalancerMetricsReporterProducer"; + public static final String DEFAULT_AUTO_BALANCER_METRICS_REPORTER_PRODUCER_CLIENT_ID = "auto_balancer_metrics_reporter_producer"; public static final long DEFAULT_AUTO_BALANCER_METRICS_REPORTER_INTERVAL_MS = TimeUnit.SECONDS.toMillis(10); public static final int DEFAULT_AUTO_BALANCER_METRICS_REPORTER_LINGER_MS = (int) TimeUnit.SECONDS.toMillis(1); public static final int DEFAULT_AUTO_BALANCER_METRICS_BATCH_SIZE = 800 * 1000; diff --git a/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfig.java b/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfig.java index 98fc70d049..22c42664e4 100644 --- a/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfig.java +++ b/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfig.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.config; diff --git a/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfigUtils.java b/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfigUtils.java index 0e263c7926..9f51245686 100644 --- a/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfigUtils.java +++ b/core/src/main/java/kafka/autobalancer/config/StaticAutoBalancerConfigUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.config; diff --git a/core/src/main/java/kafka/autobalancer/detector/AbstractAnomalyDetector.java b/core/src/main/java/kafka/autobalancer/detector/AbstractAnomalyDetector.java index a8ab8797ac..3c1e942fe8 100644 --- a/core/src/main/java/kafka/autobalancer/detector/AbstractAnomalyDetector.java +++ b/core/src/main/java/kafka/autobalancer/detector/AbstractAnomalyDetector.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.detector; diff --git a/core/src/main/java/kafka/autobalancer/detector/AnomalyDetectorImpl.java b/core/src/main/java/kafka/autobalancer/detector/AnomalyDetectorImpl.java index a474f70881..e079277887 100644 --- a/core/src/main/java/kafka/autobalancer/detector/AnomalyDetectorImpl.java +++ b/core/src/main/java/kafka/autobalancer/detector/AnomalyDetectorImpl.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.detector; @@ -44,6 +52,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; import java.util.stream.Collectors; public class AnomalyDetectorImpl extends AbstractAnomalyDetector implements LeaderChangeListener { @@ -56,6 +65,7 @@ public class AnomalyDetectorImpl extends AbstractAnomalyDetector implements Lead private final Lock configChangeLock = new ReentrantLock(); private List goalsByPriority; private Set excludedBrokers; + private Supplier> lockedNodes = Collections::emptySet; private Set excludedTopics; private long detectInterval; private long maxTolerateMetricsDelayMs; @@ -266,6 +276,10 @@ public void detect() { this.executorService.schedule(this::detect, nextExecutionDelay, TimeUnit.MILLISECONDS); } + public void lockedNodes(Supplier> lockedNodes) { + this.lockedNodes = lockedNodes; + } + private boolean isRunnable() { return this.running.get() && this.isLeader; } @@ -283,6 +297,7 @@ long detect0() throws Exception { try { detectInterval = this.detectInterval; excludedBrokers = new HashSet<>(this.excludedBrokers); + excludedBrokers.addAll(lockedNodes.get()); excludedTopics = new HashSet<>(this.excludedTopics); maxTolerateMetricsDelayMs = this.maxTolerateMetricsDelayMs; maxExecutionConcurrency = this.executionConcurrency; diff --git a/core/src/main/java/kafka/autobalancer/executor/ActionExecutorService.java b/core/src/main/java/kafka/autobalancer/executor/ActionExecutorService.java index 31ac8478b6..c278b0e0e9 100644 --- a/core/src/main/java/kafka/autobalancer/executor/ActionExecutorService.java +++ b/core/src/main/java/kafka/autobalancer/executor/ActionExecutorService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.executor; diff --git a/core/src/main/java/kafka/autobalancer/executor/ControllerActionExecutorService.java b/core/src/main/java/kafka/autobalancer/executor/ControllerActionExecutorService.java index b66a2fb55b..d54bad3bf2 100644 --- a/core/src/main/java/kafka/autobalancer/executor/ControllerActionExecutorService.java +++ b/core/src/main/java/kafka/autobalancer/executor/ControllerActionExecutorService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.executor; diff --git a/core/src/main/java/kafka/autobalancer/goals/AbstractGoal.java b/core/src/main/java/kafka/autobalancer/goals/AbstractGoal.java index a34c041630..ff01c06719 100644 --- a/core/src/main/java/kafka/autobalancer/goals/AbstractGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/AbstractGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; @@ -53,8 +61,10 @@ protected Optional tryMovePartitionOut(ActionParameters parameters) { candidateActionScores.add(Map.entry(action, score)); } } - LOGGER.debug("try move partition {} out for broker {}, all possible action score: {} on goal {}", parameters.replica.getTopicPartition(), + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("try move partition {} out for broker {}, all possible action score: {} on goal {}", parameters .replica.getTopicPartition(), parameters.srcBroker.getBrokerId(), candidateActionScores, name()); + } return getAcceptableAction(candidateActionScores); } @@ -76,8 +86,10 @@ protected Optional trySwapPartitionOut(ActionParameters parameters, candidate.getBrokerId(), candidateReplica.getTopicPartition()); double score = calculateCandidateActionScores(parameters.goalsByPriority, action, parameters.cluster, parameters.optimizedGoals, parameters.goalsByGroup); if (score > POSITIVE_ACTION_SCORE_THRESHOLD) { - LOGGER.debug("try swap partition {} out for broker {} with {}, action score: {}", parameters.replica.getTopicPartition(), + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("try swap partition {} out for broker {} with {}, action score: {}", parameters.replica.getTopicPartition(), parameters.srcBroker.getBrokerId(), candidateReplica.getTopicPartition(), score); + } return Optional.of(action); } } @@ -92,13 +104,17 @@ protected double calculateCandidateActionScores(Collection goalsByPriority for (Goal goal : goalsByPriority) { double score = goal.actionAcceptanceScore(action, cluster); if (score == NOT_ACCEPTABLE) { - LOGGER.debug("action {} is not acceptable for goal {}", action, goal); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("action {} is not acceptable for goal {}", action, goal); + } return NOT_ACCEPTABLE; } goalScoreMapByGroup.compute(goal.group(), (k, v) -> v == null ? new HashMap<>() : v).put(goal, score); } - LOGGER.debug("action {} scores on each goal: {}", action, goalScoreMapByGroup); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("action {} scores on each goal: {}", action, goalScoreMapByGroup); + } Map groupScoreMap = weightedGoalsScoreByGroup(goalScoreMapByGroup); for (Map.Entry entry : groupScoreMap.entrySet()) { String group = entry.getKey(); diff --git a/core/src/main/java/kafka/autobalancer/goals/AbstractNetworkUsageDistributionGoal.java b/core/src/main/java/kafka/autobalancer/goals/AbstractNetworkUsageDistributionGoal.java index 1d349900f5..c696756239 100644 --- a/core/src/main/java/kafka/autobalancer/goals/AbstractNetworkUsageDistributionGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/AbstractNetworkUsageDistributionGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/AbstractResourceDistributionGoal.java b/core/src/main/java/kafka/autobalancer/goals/AbstractResourceDistributionGoal.java index 055cae72b6..1ae8791ba5 100644 --- a/core/src/main/java/kafka/autobalancer/goals/AbstractResourceDistributionGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/AbstractResourceDistributionGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/AbstractResourceGoal.java b/core/src/main/java/kafka/autobalancer/goals/AbstractResourceGoal.java index 2492dcf101..90cf172719 100644 --- a/core/src/main/java/kafka/autobalancer/goals/AbstractResourceGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/AbstractResourceGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/AbstractResourceUsageDistributionGoal.java b/core/src/main/java/kafka/autobalancer/goals/AbstractResourceUsageDistributionGoal.java index 1214adea8f..1025c12f0c 100644 --- a/core/src/main/java/kafka/autobalancer/goals/AbstractResourceUsageDistributionGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/AbstractResourceUsageDistributionGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/Goal.java b/core/src/main/java/kafka/autobalancer/goals/Goal.java index faf43cad8b..e4340d53b7 100644 --- a/core/src/main/java/kafka/autobalancer/goals/Goal.java +++ b/core/src/main/java/kafka/autobalancer/goals/Goal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/GoalUtils.java b/core/src/main/java/kafka/autobalancer/goals/GoalUtils.java index 8446be6046..3eacfbb076 100644 --- a/core/src/main/java/kafka/autobalancer/goals/GoalUtils.java +++ b/core/src/main/java/kafka/autobalancer/goals/GoalUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/NetworkInUsageDistributionGoal.java b/core/src/main/java/kafka/autobalancer/goals/NetworkInUsageDistributionGoal.java index cfdbd2eb29..bb5f7153a2 100644 --- a/core/src/main/java/kafka/autobalancer/goals/NetworkInUsageDistributionGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/NetworkInUsageDistributionGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/goals/NetworkOutUsageDistributionGoal.java b/core/src/main/java/kafka/autobalancer/goals/NetworkOutUsageDistributionGoal.java index ff30d9fb94..447e850c58 100644 --- a/core/src/main/java/kafka/autobalancer/goals/NetworkOutUsageDistributionGoal.java +++ b/core/src/main/java/kafka/autobalancer/goals/NetworkOutUsageDistributionGoal.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/main/java/kafka/autobalancer/listeners/BrokerStatusListener.java b/core/src/main/java/kafka/autobalancer/listeners/BrokerStatusListener.java index b92993fd59..5e68d7d178 100644 --- a/core/src/main/java/kafka/autobalancer/listeners/BrokerStatusListener.java +++ b/core/src/main/java/kafka/autobalancer/listeners/BrokerStatusListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.listeners; diff --git a/core/src/main/java/kafka/autobalancer/listeners/ClusterStatusListenerRegistry.java b/core/src/main/java/kafka/autobalancer/listeners/ClusterStatusListenerRegistry.java index 3ec8b302e2..6bd428a14b 100644 --- a/core/src/main/java/kafka/autobalancer/listeners/ClusterStatusListenerRegistry.java +++ b/core/src/main/java/kafka/autobalancer/listeners/ClusterStatusListenerRegistry.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.listeners; diff --git a/core/src/main/java/kafka/autobalancer/listeners/LeaderChangeListener.java b/core/src/main/java/kafka/autobalancer/listeners/LeaderChangeListener.java index 971dba2dab..23f4d4b668 100644 --- a/core/src/main/java/kafka/autobalancer/listeners/LeaderChangeListener.java +++ b/core/src/main/java/kafka/autobalancer/listeners/LeaderChangeListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.listeners; diff --git a/core/src/main/java/kafka/autobalancer/listeners/TopicPartitionStatusListener.java b/core/src/main/java/kafka/autobalancer/listeners/TopicPartitionStatusListener.java index 291bf5f686..19b9863383 100644 --- a/core/src/main/java/kafka/autobalancer/listeners/TopicPartitionStatusListener.java +++ b/core/src/main/java/kafka/autobalancer/listeners/TopicPartitionStatusListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.listeners; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporter.java b/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporter.java index d96d68c6cd..e5758a05d5 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporter.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter; @@ -42,6 +50,7 @@ import org.apache.kafka.common.utils.KafkaThread; import org.apache.kafka.network.SocketServerConfigs; import org.apache.kafka.server.config.KRaftConfigs; +import org.apache.kafka.server.config.QuotaConfigs; import org.apache.kafka.server.config.ServerConfigs; import org.apache.kafka.server.metrics.KafkaYammerMetrics; @@ -229,7 +238,7 @@ public void configure(Map rawConfigs) { setIfAbsent(producerProps, ProducerConfig.CLIENT_ID_CONFIG, - reporterConfig.getString(AutoBalancerMetricsReporterConfig.AUTO_BALANCER_METRICS_REPORTER_PRODUCER_CLIENT_ID)); + QuotaConfigs.INTERNAL_CLIENT_ID_PREFIX + reporterConfig.getString(AutoBalancerMetricsReporterConfig.AUTO_BALANCER_METRICS_REPORTER_PRODUCER_CLIENT_ID)); setIfAbsent(producerProps, ProducerConfig.LINGER_MS_CONFIG, reporterConfig.getLong(AutoBalancerMetricsReporterConfig.AUTO_BALANCER_METRICS_REPORTER_LINGER_MS_CONFIG).toString()); setIfAbsent(producerProps, ProducerConfig.BATCH_SIZE_CONFIG, @@ -284,6 +293,7 @@ protected void createAutoBalancerMetricsProducer(Properties producerProps) throw }, metricsReporterCreateRetries); } + @SuppressWarnings("NPathComplexity") @Override public void run() { LOGGER.info("Starting auto balancer metrics reporter with reporting interval of {} ms.", reportingIntervalMs); @@ -291,7 +301,9 @@ public void run() { try { while (!shutdown) { long now = System.currentTimeMillis(); - LOGGER.debug("Reporting metrics for time {}.", now); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Reporting metrics for time {}.", now); + } try { if (now > lastReportingTime + reportingIntervalMs) { numMetricSendFailure = 0; @@ -342,7 +354,9 @@ public void run() { public void sendAutoBalancerMetric(AutoBalancerMetrics ccm) { ProducerRecord producerRecord = new ProducerRecord<>(Topic.AUTO_BALANCER_METRICS_TOPIC_NAME, null, ccm.time(), ccm.key(), ccm); - LOGGER.debug("Sending auto balancer metric {}.", ccm); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Sending auto balancer metric {}.", ccm); + } producer.send(producerRecord, (recordMetadata, e) -> { if (e != null) { numMetricSendFailure++; @@ -355,7 +369,9 @@ public void sendAutoBalancerMetric(AutoBalancerMetrics ccm) { } private void reportMetrics(long now) throws Exception { - LOGGER.debug("Reporting metrics."); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Reporting metrics."); + } YammerMetricProcessor.Context context = new YammerMetricProcessor.Context(now, brokerId, brokerRack, reportingIntervalMs); processMetrics(context); @@ -365,7 +381,9 @@ private void reportMetrics(long now) throws Exception { sendAutoBalancerMetric(entry.getValue()); } - LOGGER.debug("Finished reporting metrics, total metrics size: {}, merged size: {}.", interestedMetrics.size(), context.getMetricMap().size()); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Finished reporting metrics, total metrics size: {}, merged size: {}.", interestedMetrics.size(), context.getMetricMap().size()); + } } protected void checkMetricCompleteness(YammerMetricProcessor.Context context) { @@ -402,7 +420,9 @@ protected void processBrokerMetrics(YammerMetricProcessor.Context context) { protected void processYammerMetrics(YammerMetricProcessor.Context context) throws Exception { for (Map.Entry entry : interestedMetrics.entrySet()) { - LOGGER.trace("Processing yammer metric {}, scope = {}", entry.getKey(), entry.getKey().getScope()); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("Processing yammer metric {}, scope = {}", entry.getKey(), entry.getKey().getScope()); + } entry.getValue().processWith(yammerMetricProcessor, entry.getKey(), context); } Iterator> iterator = context.getMetricMap().entrySet().iterator(); @@ -429,7 +449,9 @@ protected void addMissingMetrics(YammerMetricProcessor.Context context) { protected void addMetricIfInterested(MetricName name, Metric metric) { if (isInterestedMetric(name)) { - LOGGER.debug("Added new metric {} to auto balancer metrics reporter.", name); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Added new metric {} to auto balancer metrics reporter.", name); + } interestedMetrics.put(name, metric); } } diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsUtils.java b/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsUtils.java index c4b25d0967..b4dbb52ad7 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsUtils.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/exception/UnknownVersionException.java b/core/src/main/java/kafka/autobalancer/metricsreporter/exception/UnknownVersionException.java index ef043d081f..5034c88729 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/exception/UnknownVersionException.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/exception/UnknownVersionException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.exception; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/AutoBalancerMetrics.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/AutoBalancerMetrics.java index af87ef161b..496c1751a5 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/AutoBalancerMetrics.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/AutoBalancerMetrics.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/BrokerMetrics.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/BrokerMetrics.java index 91764b9f26..45af93790b 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/BrokerMetrics.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/BrokerMetrics.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/Derivator.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/Derivator.java index d9ed4730c4..eadb2cde5a 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/Derivator.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/Derivator.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/EmptyMeter.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/EmptyMeter.java index 276df20f02..df9a7ccd32 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/EmptyMeter.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/EmptyMeter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricSerde.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricSerde.java index c3ad3a71d0..9538c34e2a 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricSerde.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricSerde.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricsUtils.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricsUtils.java index fdc095549b..bcc6c5c309 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricsUtils.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/MetricsUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/TopicPartitionMetrics.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/TopicPartitionMetrics.java index 7030e66df6..ef01f97a32 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/TopicPartitionMetrics.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/TopicPartitionMetrics.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/YammerMetricProcessor.java b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/YammerMetricProcessor.java index 427f4d1b9c..16f6f49836 100644 --- a/core/src/main/java/kafka/autobalancer/metricsreporter/metric/YammerMetricProcessor.java +++ b/core/src/main/java/kafka/autobalancer/metricsreporter/metric/YammerMetricProcessor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; @@ -41,7 +49,9 @@ public class YammerMetricProcessor implements MetricProcessor gauge, Context context) { if (MetricsUtils.isInterested(metricName)) { - LOG.trace("Processing metric {} of type Gauge.", metricName); + if (LOG.isTraceEnabled()) { + LOG.trace("Processing metric {} of type Gauge.", metricName); + } if (!(gauge.value() instanceof Number)) { throw new IllegalStateException(String.format("The value of yammer metric %s is %s, which is not a number", metricName, gauge.value())); diff --git a/core/src/main/java/kafka/autobalancer/model/AbstractInstanceUpdater.java b/core/src/main/java/kafka/autobalancer/model/AbstractInstanceUpdater.java index b30930f72d..f857342bff 100644 --- a/core/src/main/java/kafka/autobalancer/model/AbstractInstanceUpdater.java +++ b/core/src/main/java/kafka/autobalancer/model/AbstractInstanceUpdater.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/BrokerUpdater.java b/core/src/main/java/kafka/autobalancer/model/BrokerUpdater.java index d13a4194c7..4c57136278 100644 --- a/core/src/main/java/kafka/autobalancer/model/BrokerUpdater.java +++ b/core/src/main/java/kafka/autobalancer/model/BrokerUpdater.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/ClusterModel.java b/core/src/main/java/kafka/autobalancer/model/ClusterModel.java index 955fde0cec..bdc2d0030a 100644 --- a/core/src/main/java/kafka/autobalancer/model/ClusterModel.java +++ b/core/src/main/java/kafka/autobalancer/model/ClusterModel.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/ClusterModelSnapshot.java b/core/src/main/java/kafka/autobalancer/model/ClusterModelSnapshot.java index 0a1c73fea3..c1c8925239 100644 --- a/core/src/main/java/kafka/autobalancer/model/ClusterModelSnapshot.java +++ b/core/src/main/java/kafka/autobalancer/model/ClusterModelSnapshot.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/ModelUtils.java b/core/src/main/java/kafka/autobalancer/model/ModelUtils.java index b222e9db9f..8a863f22b2 100644 --- a/core/src/main/java/kafka/autobalancer/model/ModelUtils.java +++ b/core/src/main/java/kafka/autobalancer/model/ModelUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/RecordClusterModel.java b/core/src/main/java/kafka/autobalancer/model/RecordClusterModel.java index 74bff83bfd..ef144d416d 100644 --- a/core/src/main/java/kafka/autobalancer/model/RecordClusterModel.java +++ b/core/src/main/java/kafka/autobalancer/model/RecordClusterModel.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/Snapshot.java b/core/src/main/java/kafka/autobalancer/model/Snapshot.java index b15024d3cf..eff979b3d2 100644 --- a/core/src/main/java/kafka/autobalancer/model/Snapshot.java +++ b/core/src/main/java/kafka/autobalancer/model/Snapshot.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/TopicPartitionReplicaUpdater.java b/core/src/main/java/kafka/autobalancer/model/TopicPartitionReplicaUpdater.java index fca32c4f69..cdfed554d8 100644 --- a/core/src/main/java/kafka/autobalancer/model/TopicPartitionReplicaUpdater.java +++ b/core/src/main/java/kafka/autobalancer/model/TopicPartitionReplicaUpdater.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model; diff --git a/core/src/main/java/kafka/autobalancer/model/samples/SnapshottableSamples.java b/core/src/main/java/kafka/autobalancer/model/samples/SnapshottableSamples.java index 2bbc61d46a..75a8bc5ce3 100644 --- a/core/src/main/java/kafka/autobalancer/model/samples/SnapshottableSamples.java +++ b/core/src/main/java/kafka/autobalancer/model/samples/SnapshottableSamples.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model.samples; diff --git a/core/src/main/java/kafka/autobalancer/services/AbstractResumableService.java b/core/src/main/java/kafka/autobalancer/services/AbstractResumableService.java index cf8a1fc58d..024f2ffbe0 100644 --- a/core/src/main/java/kafka/autobalancer/services/AbstractResumableService.java +++ b/core/src/main/java/kafka/autobalancer/services/AbstractResumableService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.services; diff --git a/core/src/main/java/kafka/autobalancer/services/AutoBalancerService.java b/core/src/main/java/kafka/autobalancer/services/AutoBalancerService.java index b720e0363b..ea67dcb8a7 100644 --- a/core/src/main/java/kafka/autobalancer/services/AutoBalancerService.java +++ b/core/src/main/java/kafka/autobalancer/services/AutoBalancerService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.services; diff --git a/core/src/main/java/kafka/autobalancer/services/ResumableService.java b/core/src/main/java/kafka/autobalancer/services/ResumableService.java index a0f8d859bd..f2692c4c7d 100644 --- a/core/src/main/java/kafka/autobalancer/services/ResumableService.java +++ b/core/src/main/java/kafka/autobalancer/services/ResumableService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.services; diff --git a/core/src/main/java/kafka/automq/AutoMQConfig.java b/core/src/main/java/kafka/automq/AutoMQConfig.java index 6a6490cff9..bdfbe50d53 100644 --- a/core/src/main/java/kafka/automq/AutoMQConfig.java +++ b/core/src/main/java/kafka/automq/AutoMQConfig.java @@ -1,20 +1,28 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.automq; -import kafka.log.stream.s3.telemetry.exporter.ExporterConstants; import kafka.server.KafkaConfig; import org.apache.kafka.common.config.ConfigDef; +import org.apache.kafka.common.config.types.Password; import org.apache.kafka.common.utils.Utils; import com.automq.stream.s3.ByteBufAllocPolicy; @@ -28,6 +36,8 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Optional; +import java.util.concurrent.TimeUnit; import static org.apache.kafka.common.config.ConfigDef.Importance.HIGH; import static org.apache.kafka.common.config.ConfigDef.Importance.LOW; @@ -35,6 +45,7 @@ import static org.apache.kafka.common.config.ConfigDef.Type.BOOLEAN; import static org.apache.kafka.common.config.ConfigDef.Type.INT; import static org.apache.kafka.common.config.ConfigDef.Type.LONG; +import static org.apache.kafka.common.config.ConfigDef.Type.PASSWORD; import static org.apache.kafka.common.config.ConfigDef.Type.STRING; public class AutoMQConfig { @@ -44,8 +55,7 @@ public class AutoMQConfig { public static final String ELASTIC_STREAM_ENABLE_DOC = "Whether to enable AutoMQ, it has to be set to true"; public static final String ELASTIC_STREAM_ENDPOINT_CONFIG = "elasticstream.endpoint"; - public static final String ELASTIC_STREAM_ENDPOINT_DOC = "Specifies the Elastic Stream endpoint, ex. es://hostname1:port1,hostname2:port2,hostname3:port3.\n" + - "You could also PoC launch it in memory mode with endpoint memory::// or redis mode with redis://."; + public static final String ELASTIC_STREAM_ENDPOINT_DOC = "Specifies the Elastic Stream endpoint"; public static final String S3_DATA_BUCKETS_CONFIG = "s3.data.buckets"; public static final String S3_DATA_BUCKETS_DOC = "The data buckets url with format 0@s3://$bucket?region=$region. \n" + @@ -57,8 +67,7 @@ public class AutoMQConfig { public static final String S3_OPS_BUCKETS_DOC = "With the same format as s3.data.buckets"; public static final String S3_WAL_PATH_CONFIG = "s3.wal.path"; - public static final String S3_WAL_PATH_DOC = "The local WAL path for AutoMQ can be set to a block device path such as 0@file:///dev/xxx?iops=3000&iodepth=8&iobandwidth=157286400 or a filesystem file path." + - "It is recommended to use a block device for better write performance."; + public static final String S3_WAL_PATH_DOC = "The WAL path for AutoMQ, The format is '0@s3://$bucket?region=$region[&batchInterval=250][&maxBytesInBatch=8388608]'"; public static final String S3_WAL_CACHE_SIZE_CONFIG = "s3.wal.cache.size"; public static final String S3_WAL_CACHE_SIZE_DOC = "The WAL (Write-Ahead Log) cache is a FIFO (First In, First Out) queue that contains data that has not yet been uploaded to object storage, as well as data that has been uploaded but not yet evicted from the cache." + @@ -68,6 +77,10 @@ public class AutoMQConfig { public static final String S3_WAL_UPLOAD_THRESHOLD_CONFIG = "s3.wal.upload.threshold"; public static final String S3_WAL_UPLOAD_THRESHOLD_DOC = "The threshold at which WAL triggers upload to object storage. The configuration value needs to be less than s3.wal.cache.size. The larger the configuration value, the higher the data aggregation and the lower the cost of metadata storage."; + public static final String S3_WAL_UPLOAD_INTERVAL_MS_CONFIG = "s3.wal.upload.interval.ms"; + public static final String S3_WAL_UPLOAD_INTERVAL_MS_DOC = "The interval at which WAL triggers upload to object storage. -1 means only upload by size trigger"; + public static final long S3_WAL_UPLOAD_INTERVAL_MS_DEFAULT = 60000L; + public static final String S3_STREAM_SPLIT_SIZE_CONFIG = "s3.stream.object.split.size"; public static final String S3_STREAM_SPLIT_SIZE_DOC = "The S3 stream object split size threshold when upload delta WAL or compact stream set object."; @@ -100,7 +113,7 @@ public class AutoMQConfig { public static final String S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL_CONFIG = "s3.stream.set.object.compaction.interval.minutes"; public static final String S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL_DOC = "Set the interpublic static final String for Stream object compaction. The smaller this value, the smaller the scale of metadata storage, and the earlier the data can become compact. " + "However, the number of compactions that the final generated stream object goes through will increase."; - public static final int S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL = 10; // 10min + public static final int S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL = 5; // 5min public static final String S3_STREAM_SET_OBJECT_COMPACTION_CACHE_SIZE_CONFIG = "s3.stream.set.object.compaction.cache.size"; public static final String S3_STREAM_SET_OBJECT_COMPACTION_CACHE_SIZE_DOC = "The size of memory is available during the Stream object compaction process. The larger this value, the lower the cost of API calls."; @@ -140,7 +153,7 @@ public class AutoMQConfig { public static final String S3_NETWORK_BASELINE_BANDWIDTH_CONFIG = "s3.network.baseline.bandwidth"; public static final String S3_NETWORK_BASELINE_BANDWIDTH_DOC = "The total available bandwidth for object storage requests. This is used to prevent stream set object compaction and catch-up read from monopolizing normal read and write traffic. Produce and Consume will also separately consume traffic in and traffic out. " + "For example, suppose this value is set to 100MB/s, and the normal read and write traffic is 80MB/s, then the available traffic for stream set object compaction is 20MB/s."; - public static final long S3_NETWORK_BASELINE_BANDWIDTH = 100 * 1024 * 1024; // 100MB/s + public static final long S3_NETWORK_BASELINE_BANDWIDTH = 1024 * 1024 * 1024; // 1GBps public static final String S3_NETWORK_REFILL_PERIOD_MS_CONFIG = "s3.network.refill.period.ms"; public static final String S3_NETWORK_REFILL_PERIOD_MS_DOC = "The network bandwidth token refill period in milliseconds."; @@ -171,6 +184,21 @@ public class AutoMQConfig { public static final String CLUSTER_ID_CONFIG = "cluster.id"; public static final String CLUSTER_ID_DOC = "If the cluster.id is set, Kafka will auto format the storage."; + public static final String S3_BACK_PRESSURE_ENABLED_CONFIG = "automq.backpressure.enabled"; + public static final String S3_BACK_PRESSURE_ENABLED_DOC = "Whether back pressure is enabled"; + public static final boolean S3_BACK_PRESSURE_ENABLED_DEFAULT = true; + + public static final String S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG = "automq.backpressure.cooldown.ms"; + public static final String S3_BACK_PRESSURE_COOLDOWN_MS_DOC = "The cooldown time in milliseconds to wait between two regulator actions"; + public static final long S3_BACK_PRESSURE_COOLDOWN_MS_DEFAULT = TimeUnit.SECONDS.toMillis(15); + + public static final String TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG = "automq.table.topic.schema.registry.url"; + private static final String TABLE_TOPIC_SCHEMA_REGISTRY_URL_DOC = "The schema registry url for table topic"; + + public static final String ZONE_ROUTER_CHANNELS_CONFIG = "automq.zonerouter.channels"; + public static final String ZONE_ROUTER_CHANNELS_DOC = "The channels to use for cross zone router. Currently it only support object storage channel." + + " The format is '0@s3://$bucket?region=$region[&batchInterval=250][&maxBytesInBatch=8388608]'"; + // Deprecated config start public static final String S3_ENDPOINT_CONFIG = "s3.endpoint"; public static final String S3_ENDPOINT_DOC = "[DEPRECATED]please use s3.data.buckets. The object storage endpoint, ex. https://s3.us-east-1.amazonaws.com."; @@ -221,6 +249,12 @@ public class AutoMQConfig { public static final String S3_TELEMETRY_OPS_ENABLED_CONFIG = "s3.telemetry.ops.enabled"; public static final String S3_TELEMETRY_OPS_ENABLED_DOC = "[DEPRECATED] use s3.telemetry.metrics.uri instead."; + + private static final String TELEMETRY_EXPORTER_TYPE_OTLP = "otlp"; + private static final String TELEMETRY_EXPORTER_TYPE_PROMETHEUS = "prometheus"; + private static final String TELEMETRY_EXPORTER_TYPE_OPS = "ops"; + public static final String URI_DELIMITER = "://?"; + // Deprecated config end public static void define(ConfigDef configDef) { @@ -231,13 +265,14 @@ public static void define(ConfigDef configDef) { .define(AutoMQConfig.S3_WAL_PATH_CONFIG, STRING, null, HIGH, AutoMQConfig.S3_WAL_PATH_DOC) .define(AutoMQConfig.S3_WAL_CACHE_SIZE_CONFIG, LONG, -1L, MEDIUM, AutoMQConfig.S3_WAL_CACHE_SIZE_DOC) .define(AutoMQConfig.S3_WAL_UPLOAD_THRESHOLD_CONFIG, LONG, -1L, MEDIUM, AutoMQConfig.S3_WAL_UPLOAD_THRESHOLD_DOC) + .define(AutoMQConfig.S3_WAL_UPLOAD_INTERVAL_MS_CONFIG, LONG, S3_WAL_UPLOAD_INTERVAL_MS_DEFAULT, LOW, AutoMQConfig.S3_WAL_UPLOAD_INTERVAL_MS_DOC) .define(AutoMQConfig.S3_STREAM_SPLIT_SIZE_CONFIG, INT, 8388608, MEDIUM, AutoMQConfig.S3_STREAM_SPLIT_SIZE_DOC) .define(AutoMQConfig.S3_OBJECT_BLOCK_SIZE_CONFIG, INT, 524288, MEDIUM, AutoMQConfig.S3_OBJECT_BLOCK_SIZE_DOC) .define(AutoMQConfig.S3_OBJECT_PART_SIZE_CONFIG, INT, 16777216, MEDIUM, AutoMQConfig.S3_OBJECT_PART_SIZE_DOC) .define(AutoMQConfig.S3_BLOCK_CACHE_SIZE_CONFIG, LONG, -1L, MEDIUM, AutoMQConfig.S3_BLOCK_CACHE_SIZE_DOC) .define(AutoMQConfig.S3_STREAM_ALLOCATOR_POLICY_CONFIG, STRING, ByteBufAllocPolicy.POOLED_HEAP.name(), MEDIUM, AutoMQConfig.S3_STREAM_ALLOCATOR_POLICY_DOC) .define(AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_INTERVAL_MINUTES_CONFIG, INT, 30, MEDIUM, AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_INTERVAL_MINUTES_DOC) - .define(AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_MAX_SIZE_BYTES_CONFIG, LONG, 1073741824L, MEDIUM, AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_MAX_SIZE_BYTES_DOC) + .define(AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_MAX_SIZE_BYTES_CONFIG, LONG, 10737418240L, MEDIUM, AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_MAX_SIZE_BYTES_DOC) .define(AutoMQConfig.S3_CONTROLLER_REQUEST_RETRY_MAX_COUNT_CONFIG, INT, Integer.MAX_VALUE, MEDIUM, AutoMQConfig.S3_CONTROLLER_REQUEST_RETRY_MAX_COUNT_DOC) .define(AutoMQConfig.S3_CONTROLLER_REQUEST_RETRY_BASE_DELAY_MS_CONFIG, LONG, 500, MEDIUM, AutoMQConfig.S3_CONTROLLER_REQUEST_RETRY_BASE_DELAY_MS_DOC) .define(AutoMQConfig.S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL_CONFIG, INT, S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL, MEDIUM, AutoMQConfig.S3_STREAM_SET_OBJECT_COMPACTION_INTERVAL_DOC) @@ -253,8 +288,11 @@ public static void define(ConfigDef configDef) { .define(AutoMQConfig.S3_NETWORK_REFILL_PERIOD_MS_CONFIG, INT, S3_REFILL_PERIOD_MS, MEDIUM, AutoMQConfig.S3_NETWORK_REFILL_PERIOD_MS_DOC) .define(AutoMQConfig.S3_TELEMETRY_METRICS_LEVEL_CONFIG, STRING, "INFO", MEDIUM, AutoMQConfig.S3_TELEMETRY_METRICS_LEVEL_DOC) .define(AutoMQConfig.S3_TELEMETRY_EXPORTER_REPORT_INTERVAL_MS_CONFIG, INT, S3_METRICS_EXPORTER_REPORT_INTERVAL_MS, MEDIUM, AutoMQConfig.S3_TELEMETRY_EXPORTER_REPORT_INTERVAL_MS_DOC) - .define(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG, STRING, null, HIGH, AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_DOC) + .define(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG, PASSWORD, null, HIGH, AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_DOC) .define(AutoMQConfig.S3_TELEMETRY_METRICS_BASE_LABELS_CONFIG, STRING, null, MEDIUM, AutoMQConfig.S3_TELEMETRY_METRICS_BASE_LABELS_DOC) + .define(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, BOOLEAN, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_DEFAULT, MEDIUM, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_DOC) + .define(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG, LONG, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_DEFAULT, MEDIUM, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_DOC) + .define(AutoMQConfig.ZONE_ROUTER_CHANNELS_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.HIGH, AutoMQConfig.ZONE_ROUTER_CHANNELS_DOC) // Deprecated config start .define(AutoMQConfig.S3_ENDPOINT_CONFIG, STRING, null, HIGH, AutoMQConfig.S3_ENDPOINT_DOC) .define(AutoMQConfig.S3_REGION_CONFIG, STRING, null, HIGH, AutoMQConfig.S3_REGION_DOC) @@ -271,14 +309,18 @@ public static void define(ConfigDef configDef) { .define(AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_PROTOCOL_CONFIG, STRING, S3_EXPORTER_OTLPPROTOCOL, MEDIUM, AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_PROTOCOL_DOC) .define(AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_COMPRESSION_ENABLE_CONFIG, BOOLEAN, false, MEDIUM, AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_COMPRESSION_ENABLE_DOC) .define(AutoMQConfig.S3_METRICS_EXPORTER_PROM_HOST_CONFIG, STRING, "localhost", MEDIUM, AutoMQConfig.S3_METRICS_EXPORTER_PROM_HOST_DOC) - .define(AutoMQConfig.S3_METRICS_EXPORTER_PROM_PORT_CONFIG, INT, 9090, MEDIUM, AutoMQConfig.S3_METRICS_EXPORTER_PROM_PORT_DOC); + .define(AutoMQConfig.S3_METRICS_EXPORTER_PROM_PORT_CONFIG, INT, 9090, MEDIUM, AutoMQConfig.S3_METRICS_EXPORTER_PROM_PORT_DOC) + .define(AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG, STRING, null, MEDIUM, AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_DOC); } + private final long nodeEpoch = System.currentTimeMillis(); private List dataBuckets; private List opsBuckets; private String walConfig; private String metricsExporterURI; private List> baseLabels; + @SuppressWarnings("OptionalUsedAsFieldOrParameterType") + private Optional> zoneRouterChannels; public AutoMQConfig setup(KafkaConfig config) { dataBuckets = genDataBuckets(config); @@ -286,9 +328,14 @@ public AutoMQConfig setup(KafkaConfig config) { walConfig = genWALConfig(config); metricsExporterURI = genMetricsExporterURI(config); baseLabels = parseBaseLabels(config); + zoneRouterChannels = genZoneRouterChannels(config); return this; } + public long nodeEpoch() { + return nodeEpoch; + } + public List dataBuckets() { return dataBuckets; } @@ -309,6 +356,10 @@ public List> baseLabels() { return baseLabels; } + public Optional> zoneRouterChannels() { + return zoneRouterChannels; + } + private static List genDataBuckets(KafkaConfig config) { String dataBuckets = config.getString(S3_DATA_BUCKETS_CONFIG); String oldEndpoint = config.getString(S3_ENDPOINT_CONFIG); @@ -351,11 +402,12 @@ private static String genWALConfig(KafkaConfig config) { } private static String genMetricsExporterURI(KafkaConfig config) { - String uri = config.getString(S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG); + Password pwd = config.getPassword(S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG); + String uri = pwd == null ? null : pwd.value(); if (uri == null) { uri = buildMetrixExporterURIWithOldConfigs(config); } - if (!uri.contains(ExporterConstants.OPS_TYPE)) { + if (!uri.contains(TELEMETRY_EXPORTER_TYPE_OPS)) { uri += "," + buildOpsExporterURI(); } return uri; @@ -372,10 +424,10 @@ private static String buildMetrixExporterURIWithOldConfigs(KafkaConfig kafkaConf for (String exporterType : exporterTypeArray) { exporterType = exporterType.trim(); switch (exporterType) { - case ExporterConstants.OTLP_TYPE: + case TELEMETRY_EXPORTER_TYPE_OTLP: exportedUris.add(buildOTLPExporterURI(kafkaConfig)); break; - case ExporterConstants.PROMETHEUS_TYPE: + case TELEMETRY_EXPORTER_TYPE_PROMETHEUS: exportedUris.add(buildPrometheusExporterURI(kafkaConfig)); break; default: @@ -393,26 +445,31 @@ private static String buildMetrixExporterURIWithOldConfigs(KafkaConfig kafkaConf } private static String buildOTLPExporterURI(KafkaConfig kafkaConfig) { + String endpoint = kafkaConfig.getString(S3_TELEMETRY_EXPORTER_OTLP_ENDPOINT_CONFIG); + if (StringUtils.isBlank(endpoint)) { + return ""; + } StringBuilder uriBuilder = new StringBuilder() - .append(ExporterConstants.OTLP_TYPE) - .append(ExporterConstants.URI_DELIMITER) - .append(ExporterConstants.ENDPOINT).append("=").append(kafkaConfig.getString(S3_TELEMETRY_EXPORTER_OTLP_ENDPOINT_CONFIG)) - .append("&") - .append(ExporterConstants.PROTOCOL).append("=").append(kafkaConfig.getString(S3_TELEMETRY_EXPORTER_OTLP_PROTOCOL_CONFIG)); + .append(TELEMETRY_EXPORTER_TYPE_OTLP) + .append("://?endpoint=").append(endpoint); + String protocol = kafkaConfig.getString(S3_TELEMETRY_EXPORTER_OTLP_PROTOCOL_CONFIG); + if (StringUtils.isNotBlank(protocol)) { + uriBuilder.append("&protocol=").append(protocol); + } if (kafkaConfig.getBoolean(S3_TELEMETRY_EXPORTER_OTLP_COMPRESSION_ENABLE_CONFIG)) { - uriBuilder.append("&").append(ExporterConstants.COMPRESSION).append("=").append("gzip"); + uriBuilder.append("&compression=gzip"); } return uriBuilder.toString(); } private static String buildPrometheusExporterURI(KafkaConfig kafkaConfig) { - return ExporterConstants.PROMETHEUS_TYPE + ExporterConstants.URI_DELIMITER + - ExporterConstants.HOST + "=" + kafkaConfig.getString(S3_METRICS_EXPORTER_PROM_HOST_CONFIG) + "&" + - ExporterConstants.PORT + "=" + kafkaConfig.getInt(S3_METRICS_EXPORTER_PROM_PORT_CONFIG); + return TELEMETRY_EXPORTER_TYPE_PROMETHEUS + URI_DELIMITER + + "host" + "=" + kafkaConfig.getString(S3_METRICS_EXPORTER_PROM_HOST_CONFIG) + "&" + + "port" + "=" + kafkaConfig.getInt(S3_METRICS_EXPORTER_PROM_PORT_CONFIG); } private static String buildOpsExporterURI() { - return ExporterConstants.OPS_TYPE + ExporterConstants.URI_DELIMITER; + return TELEMETRY_EXPORTER_TYPE_OPS + URI_DELIMITER; } private static List> parseBaseLabels(KafkaConfig config) { @@ -431,4 +488,18 @@ private static List> parseBaseLabels(KafkaConfig config) { } return labels; } + + + private static Optional> genZoneRouterChannels(KafkaConfig config) { + String str = config.getString(ZONE_ROUTER_CHANNELS_CONFIG); + if (StringUtils.isBlank(str)) { + return Optional.empty(); + } + List buckets = BucketURI.parseBuckets(str); + if (buckets.isEmpty()) { + return Optional.empty(); + } else { + return Optional.of(buckets); + } + } } diff --git a/core/src/main/java/kafka/automq/StorageUtil.java b/core/src/main/java/kafka/automq/StorageUtil.java index dec89ff40e..8031129d6b 100644 --- a/core/src/main/java/kafka/automq/StorageUtil.java +++ b/core/src/main/java/kafka/automq/StorageUtil.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.automq; @@ -58,4 +66,4 @@ private static String persistConfig(Properties serverProps, String basePath) { } } -} \ No newline at end of file +} diff --git a/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java b/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java new file mode 100644 index 0000000000..d8c7c946ce --- /dev/null +++ b/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java @@ -0,0 +1,95 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +import kafka.automq.AutoMQConfig; +import kafka.server.KafkaConfig; + +import org.apache.kafka.common.config.ConfigException; +import org.apache.kafka.common.utils.ConfigUtils; + +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +public class BackPressureConfig { + + public static final Set RECONFIGURABLE_CONFIGS = Set.of( + AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, + AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG + ); + + private volatile boolean enabled; + /** + * The cooldown time in milliseconds to wait between two regulator actions. + */ + private long cooldownMs; + + public static BackPressureConfig from(KafkaConfig config) { + return new BackPressureConfig(config.s3BackPressureEnabled(), config.s3BackPressureCooldownMs()); + } + + public static BackPressureConfig from(Map raw) { + Map configs = new HashMap<>(raw); + return new BackPressureConfig( + ConfigUtils.getBoolean(configs, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG), + ConfigUtils.getLong(configs, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG) + ); + } + + public BackPressureConfig(boolean enabled, long cooldownMs) { + this.enabled = enabled; + this.cooldownMs = cooldownMs; + } + + public static void validate(Map raw) throws ConfigException { + Map configs = new HashMap<>(raw); + if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG)) { + ConfigUtils.getBoolean(configs, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG); + } + if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)) { + validateCooldownMs(ConfigUtils.getLong(configs, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)); + } + } + + public static void validateCooldownMs(long cooldownMs) throws ConfigException { + if (cooldownMs < 0) { + throw new ConfigException(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG, cooldownMs, "The cooldown time must be non-negative."); + } + } + + public void update(Map raw) { + Map configs = new HashMap<>(raw); + if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG)) { + this.enabled = ConfigUtils.getBoolean(configs, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG); + } + if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)) { + this.cooldownMs = ConfigUtils.getLong(configs, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG); + } + } + + public boolean enabled() { + return enabled; + } + + public long cooldownMs() { + return cooldownMs; + } +} diff --git a/core/src/main/java/kafka/automq/backpressure/BackPressureManager.java b/core/src/main/java/kafka/automq/backpressure/BackPressureManager.java new file mode 100644 index 0000000000..5f3840d25f --- /dev/null +++ b/core/src/main/java/kafka/automq/backpressure/BackPressureManager.java @@ -0,0 +1,45 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +import org.apache.kafka.common.Reconfigurable; + +/** + * It checks the {@link LoadLevel} of the system and takes actions based on the load level + * to prevent the system from being overwhelmed. + */ +public interface BackPressureManager extends Reconfigurable { + + /** + * Start the back pressure manager. + */ + void start(); + + /** + * Register a checker to check the load level of the system. + * Note: It should be called between {@link #start()} and {@link #shutdown()}. + */ + void registerChecker(Checker checker); + + /** + * Shutdown the back pressure manager, and release all resources. + */ + void shutdown(); +} diff --git a/core/src/main/java/kafka/automq/backpressure/Checker.java b/core/src/main/java/kafka/automq/backpressure/Checker.java new file mode 100644 index 0000000000..babfc8be42 --- /dev/null +++ b/core/src/main/java/kafka/automq/backpressure/Checker.java @@ -0,0 +1,41 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +/** + * A checker to check the load level of the system periodically. + */ +public interface Checker { + + /** + * The source of the checker, which should be unique to identify the checker. + */ + String source(); + + /** + * Check the load level of the system. + */ + LoadLevel check(); + + /** + * The interval in milliseconds to check the load level of the system. + */ + long intervalMs(); +} diff --git a/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java b/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java new file mode 100644 index 0000000000..3abe0acf19 --- /dev/null +++ b/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java @@ -0,0 +1,179 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +import org.apache.kafka.common.config.ConfigException; +import org.apache.kafka.server.metrics.s3stream.S3StreamKafkaMetricsManager; + +import com.automq.stream.utils.ThreadUtils; +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import static kafka.automq.backpressure.BackPressureConfig.RECONFIGURABLE_CONFIGS; + +public class DefaultBackPressureManager implements BackPressureManager { + + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultBackPressureManager.class); + + private final BackPressureConfig config; + private final Regulator regulator; + + /** + * The scheduler to schedule the checker periodically. + * Package-private for testing. + */ + ScheduledExecutorService checkerScheduler; + /** + * The map to store the source and the most recent load level from the checker. + * Note: It should only be accessed in the {@link #checkerScheduler} thread. + */ + private final Map loadLevels = new HashMap<>(); + /** + * The last time to trigger the regulator. + * Note: It should only be accessed in the {@link #checkerScheduler} thread. + */ + private long lastRegulateTime = System.currentTimeMillis(); + /** + * The last load level to trigger the regulator. + * Only used for logging and monitoring. + */ + private LoadLevel lastRegulateLevel = LoadLevel.NORMAL; + /** + * The current state metrics of the system. + * Only used for monitoring. + * + * @see S3StreamKafkaMetricsManager#setBackPressureStateSupplier + */ + private final Map stateMetrics = new HashMap<>(LoadLevel.values().length); + + public DefaultBackPressureManager(BackPressureConfig config, Regulator regulator) { + this.config = config; + this.regulator = regulator; + } + + @Override + public void start() { + this.checkerScheduler = Threads.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("back-pressure-checker-%d", false), LOGGER, true, false); + S3StreamKafkaMetricsManager.setBackPressureStateSupplier(this::stateMetrics); + } + + @Override + public void registerChecker(Checker checker) { + checkerScheduler.scheduleWithFixedDelay(() -> { + loadLevels.put(checker.source(), checker.check()); + maybeRegulate(); + }, 0, checker.intervalMs(), TimeUnit.MILLISECONDS); + } + + @Override + public void shutdown() { + ThreadUtils.shutdownExecutor(checkerScheduler, 1, TimeUnit.SECONDS); + } + + private void maybeRegulate() { + if (!config.enabled()) { + return; + } + maybeRegulate(false); + } + + /** + * Regulate the system if the cooldown time has passed. + * + * @param isInternal True if it is an internal call, which means it should not schedule the next regulate action. + */ + private void maybeRegulate(boolean isInternal) { + LoadLevel loadLevel = currentLoadLevel(); + long now = System.currentTimeMillis(); + long timeElapsed = now - lastRegulateTime; + + if (timeElapsed < config.cooldownMs()) { + // Skip regulating if the cooldown time has not passed. + if (!isInternal) { + // Schedule the next regulate action if it is not an internal call. + checkerScheduler.schedule(() -> maybeRegulate(true), config.cooldownMs() - timeElapsed, TimeUnit.MILLISECONDS); + } + return; + } + regulate(loadLevel, now); + } + + /** + * Get the current load level of the system, which is, the maximum load level from all checkers. + */ + private LoadLevel currentLoadLevel() { + return loadLevels.values().stream() + .max(LoadLevel::compareTo) + .orElse(LoadLevel.NORMAL); + } + + private void regulate(LoadLevel loadLevel, long now) { + if (LoadLevel.NORMAL.equals(loadLevel)) { + if (!LoadLevel.NORMAL.equals(lastRegulateLevel)) { + LOGGER.info("The system is back to a normal state, checkers: {}", loadLevels); + } + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("The system is in a normal state, checkers: {}", loadLevels); + } + } else { + LOGGER.info("The system is in a {} state, checkers: {}", loadLevel, loadLevels); + } + + loadLevel.regulate(regulator); + lastRegulateTime = now; + lastRegulateLevel = loadLevel; + } + + private Map stateMetrics() { + LoadLevel current = currentLoadLevel(); + for (LoadLevel level : LoadLevel.values()) { + int value = level.equals(current) ? current.ordinal() : -1; + stateMetrics.put(level.name(), value); + } + return stateMetrics; + } + + @Override + public Set reconfigurableConfigs() { + return RECONFIGURABLE_CONFIGS; + } + + @Override + public void validateReconfiguration(Map configs) throws ConfigException { + BackPressureConfig.validate(configs); + } + + @Override + public void reconfigure(Map configs) { + config.update(configs); + } + + @Override + public void configure(Map configs) { + } +} diff --git a/core/src/main/java/kafka/automq/backpressure/LoadLevel.java b/core/src/main/java/kafka/automq/backpressure/LoadLevel.java new file mode 100644 index 0000000000..d3303042eb --- /dev/null +++ b/core/src/main/java/kafka/automq/backpressure/LoadLevel.java @@ -0,0 +1,51 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +/** + * Represents the load level of the system. + * {@link BackPressureManager} will take actions based on the load level. + * Note: It MUST be ordered by the severity. + */ +public enum LoadLevel { + /** + * The system is in a normal state. + */ + NORMAL { + @Override + public void regulate(Regulator regulator) { + regulator.increase(); + } + }, + /** + * The system is in a high load state, and some actions should be taken to reduce the load. + */ + HIGH { + @Override + public void regulate(Regulator regulator) { + regulator.decrease(); + } + }; + + /** + * Take actions based on the load level. + */ + public abstract void regulate(Regulator regulator); +} diff --git a/core/src/main/java/kafka/automq/backpressure/Regulator.java b/core/src/main/java/kafka/automq/backpressure/Regulator.java new file mode 100644 index 0000000000..9b0f591991 --- /dev/null +++ b/core/src/main/java/kafka/automq/backpressure/Regulator.java @@ -0,0 +1,39 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +/** + * The Regulator class is responsible for controlling and limiting the rate of external requests. + * It provides methods to increase, decrease, and minimize the flow of incoming requests. + */ +public interface Regulator { + + /** + * Increase the rate of incoming requests. + * If the rate is already at the maximum, this method does nothing. + */ + void increase(); + + /** + * Decrease the rate of incoming requests. + * If the rate is already at the minimum, this method does nothing. + */ + void decrease(); +} diff --git a/core/src/main/java/kafka/automq/controller/DefaultQuorumControllerExtension.java b/core/src/main/java/kafka/automq/controller/DefaultQuorumControllerExtension.java new file mode 100644 index 0000000000..ad3d3f98dd --- /dev/null +++ b/core/src/main/java/kafka/automq/controller/DefaultQuorumControllerExtension.java @@ -0,0 +1,57 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.controller; + +import kafka.automq.failover.FailoverControlManager; + +import org.apache.kafka.common.metadata.KVRecord; +import org.apache.kafka.common.metadata.MetadataRecordType; +import org.apache.kafka.common.protocol.ApiMessage; +import org.apache.kafka.controller.QuorumController; +import org.apache.kafka.controller.QuorumControllerExtension; +import org.apache.kafka.raft.OffsetAndEpoch; + +import java.util.Objects; +import java.util.Optional; + +public class DefaultQuorumControllerExtension implements QuorumControllerExtension { + private final FailoverControlManager failoverControlManager; + + public DefaultQuorumControllerExtension(QuorumController controller) { + this.failoverControlManager = new FailoverControlManager( + controller.snapshotRegistry(), + controller, + controller.clusterControl(), + controller.nodeControlManager(), + controller.streamControlManager() + ); + } + + @Override + public boolean replay(MetadataRecordType type, ApiMessage message, Optional snapshotId, + long batchLastOffset) { + if (Objects.requireNonNull(type) == MetadataRecordType.KVRECORD) { + failoverControlManager.replay((KVRecord) message); + } else { + return false; + } + return true; + } +} diff --git a/core/src/main/java/kafka/automq/failover/DefaultFailedNode.java b/core/src/main/java/kafka/automq/failover/DefaultFailedNode.java new file mode 100644 index 0000000000..9a1e06ec27 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/DefaultFailedNode.java @@ -0,0 +1,23 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +public record DefaultFailedNode(int id, long epoch) implements FailedNode { +} diff --git a/core/src/main/java/kafka/automq/failover/DefaultFailedWal.java b/core/src/main/java/kafka/automq/failover/DefaultFailedWal.java new file mode 100644 index 0000000000..1ede39997e --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/DefaultFailedWal.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +import java.util.Objects; + +public final class DefaultFailedWal implements FailedWal { + private final NodeRuntimeMetadata nodeMetadata; + + public DefaultFailedWal(NodeRuntimeMetadata nodeMetadata) { + this.nodeMetadata = nodeMetadata; + } + + public NodeRuntimeMetadata nodeMetadata() { + return nodeMetadata; + } + + @Override + public boolean equals(Object obj) { + if (obj == this) + return true; + if (obj == null || obj.getClass() != this.getClass()) + return false; + var that = (DefaultFailedWal) obj; + return Objects.equals(this.nodeMetadata, that.nodeMetadata); + } + + @Override + public int hashCode() { + return Objects.hash(nodeMetadata); + } + + @Override + public String toString() { + return "FailedWalV1[" + + "nodeMetadata=" + nodeMetadata + ']'; + } + + @Override + public FailoverContext toFailoverContext(int target) { + return new FailoverContext(nodeMetadata.id(), nodeMetadata.epoch(), target, nodeMetadata.walConfigs()); + } + + @Override + public FailedNode node() { + return FailedNode.from(nodeMetadata); + } + + public static FailedWal from(NodeRuntimeMetadata failedNode) { + return new DefaultFailedWal(failedNode); + } +} diff --git a/core/src/main/java/kafka/automq/failover/FailedNode.java b/core/src/main/java/kafka/automq/failover/FailedNode.java new file mode 100644 index 0000000000..a6d83aef95 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/FailedNode.java @@ -0,0 +1,36 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +public interface FailedNode { + + /** + * The node id of the failed node. + */ + int id(); + + static FailedNode from(NodeRuntimeMetadata node) { + return new DefaultFailedNode(node.id(), node.epoch()); + } + + static FailedNode from(FailoverContext context) { + return new DefaultFailedNode(context.getNodeId(), context.getNodeEpoch()); + } +} diff --git a/core/src/main/java/kafka/automq/failover/FailedWal.java b/core/src/main/java/kafka/automq/failover/FailedWal.java new file mode 100644 index 0000000000..bec42f2ab8 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/FailedWal.java @@ -0,0 +1,37 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +public interface FailedWal { + + /** + * Convert to a failover context. + */ + FailoverContext toFailoverContext(int target); + + /** + * The node of the failed WAL belongs to. + */ + FailedNode node(); + + default int nodeId() { + return node().id(); + } +} diff --git a/core/src/main/java/kafka/automq/failover/FailoverConstants.java b/core/src/main/java/kafka/automq/failover/FailoverConstants.java new file mode 100644 index 0000000000..f389a76247 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/FailoverConstants.java @@ -0,0 +1,25 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +public class FailoverConstants { + public static final String FAILOVER_KEY = "__a.failover"; + +} diff --git a/core/src/main/java/kafka/automq/failover/FailoverContext.java b/core/src/main/java/kafka/automq/failover/FailoverContext.java new file mode 100644 index 0000000000..a9f0bb6399 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/FailoverContext.java @@ -0,0 +1,99 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +@JsonIgnoreProperties(ignoreUnknown = true) +public class FailoverContext { + /** + * Failed node id + */ + @JsonProperty("n") + private int nodeId; + + /** + * Failover target node id + * + * @since failover v0 + */ + @JsonProperty("t") + private int target; + + /** + * Failed node epoch + * + * @since failover v1 + */ + @JsonInclude(JsonInclude.Include.NON_DEFAULT) + @JsonProperty(value = "e", defaultValue = "0") + private long nodeEpoch; + + /** + * WAL configs for failover + */ + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("c") + private String kraftWalConfigs; + + // for json deserialize + public FailoverContext() {} + + public FailoverContext(int nodeId, long nodeEpoch, int target, String kraftWalConfigs) { + this.nodeId = nodeId; + this.nodeEpoch = nodeEpoch; + this.target = target; + this.kraftWalConfigs = kraftWalConfigs; + } + + @JsonIgnore + public FailedNode getFailedNode() { + return FailedNode.from(this); + } + + public int getNodeId() { + return nodeId; + } + + public int getTarget() { + return target; + } + + public long getNodeEpoch() { + return nodeEpoch; + } + + public String getKraftWalConfigs() { + return kraftWalConfigs; + } + + @Override + public String toString() { + return "FailoverContext{" + + "nodeId=" + nodeId + + ", target=" + target + + ", nodeEpoch=" + nodeEpoch + + ", kraftWalConfigs=" + kraftWalConfigs + + '}'; + } +} diff --git a/core/src/main/java/kafka/automq/failover/FailoverControlManager.java b/core/src/main/java/kafka/automq/failover/FailoverControlManager.java new file mode 100644 index 0000000000..13b1ee4056 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/FailoverControlManager.java @@ -0,0 +1,270 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +import kafka.automq.utils.JsonUtils; + +import org.apache.kafka.common.metadata.KVRecord; +import org.apache.kafka.common.utils.ThreadUtils; +import org.apache.kafka.controller.ClusterControlManager; +import org.apache.kafka.controller.ControllerResult; +import org.apache.kafka.controller.QuorumController; +import org.apache.kafka.controller.stream.NodeControlManager; +import org.apache.kafka.controller.stream.StreamControlManager; +import org.apache.kafka.metadata.BrokerRegistration; +import org.apache.kafka.server.common.ApiMessageAndVersion; +import org.apache.kafka.timeline.SnapshotRegistry; +import org.apache.kafka.timeline.TimelineObject; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalLong; +import java.util.Set; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import static kafka.automq.failover.FailoverConstants.FAILOVER_KEY; + +@SuppressWarnings({"NPathComplexity"}) +public class FailoverControlManager implements AutoCloseable { + private static final Logger LOGGER = LoggerFactory.getLogger(FailoverControlManager.class); + private static final int MAX_FAILOVER_COUNT_IN_TARGET_NODE = 1; + /** + * failover contexts: failedNode -> context + */ + private final TimelineObject> failoverContexts; + private final QuorumController quorumController; + private final ClusterControlManager clusterControlManager; + private final NodeControlManager nodeControlManager; + private final StreamControlManager streamControlManager; + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("failover-controller", true)); + + public FailoverControlManager( + SnapshotRegistry registry, + QuorumController quorumController, + ClusterControlManager clusterControlManager, + NodeControlManager nodeControlManager, + StreamControlManager streamControlManager + ) { + this.failoverContexts = new TimelineObject<>(registry, Collections.emptyMap()); + this.quorumController = quorumController; + this.clusterControlManager = clusterControlManager; + this.nodeControlManager = nodeControlManager; + this.streamControlManager = streamControlManager; + this.scheduler.scheduleWithFixedDelay(this::runFailoverTask, 10, 10, TimeUnit.SECONDS); + } + + void runFailoverTask() { + if (!quorumController.isActive()) { + return; + } + try { + this.quorumController.appendWriteEvent("failover", OptionalLong.empty(), this::failover).get(); + } catch (Throwable e) { + LOGGER.error("run failover task failed", e); + } + } + + private ControllerResult failover() { + List allNodes = allNodes(); + List failedWalList = getFailedWal(allNodes); + + Map failoverContexts = this.failoverContexts.get(); + + List newFailover = addNewFailover(failoverContexts, failedWalList); + List completed = delCompleted(failoverContexts, failedWalList); + List reset = resetNotAlive(failoverContexts); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("[FAILOVER_TASK], new={}, completed={}, reset={}", newFailover, completed, reset); + } + + failoverContexts = new HashMap<>(failoverContexts); + for (FailedNode node : completed) { + FailoverContext ctx = failoverContexts.remove(node); + LOGGER.info("[FAILOVER_COMPLETE],context={}", ctx); + } + for (FailedNode node : reset) { + // remove reset nodes, and they will be added again in the next round + FailoverContext ctx = failoverContexts.remove(node); + LOGGER.info("[RESET_NOT_ALIVE_FAILOVER],context={}", ctx); + } + int maxInflight = 1; + List excess = failoverContexts.keySet().stream() + .skip(maxInflight) + .collect(Collectors.toList()); + for (FailedNode node : excess) { + FailoverContext ctx = failoverContexts.remove(node); + LOGGER.info("[REMOVE_EXCESS_FAILOVER],context={}", ctx); + } + for (FailoverContext ctx : newFailover) { + if (failoverContexts.size() < maxInflight) { + failoverContexts.put(ctx.getFailedNode(), ctx); + LOGGER.info("[ADD_NEW_FAILOVER],context={}", ctx); + } else { + LOGGER.info("[PENDING_FAILOVER],context={}", ctx); + } + } + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("[INFLIGHT_FAILOVER],contexts={}", failoverContexts); + } + if (newFailover.isEmpty() && completed.isEmpty() && reset.isEmpty() && excess.isEmpty()) { + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("No changes in failover contexts"); + } + return ControllerResult.of(Collections.emptyList(), null); + } + String encoded = JsonUtils.encode(failoverContexts.values()); + ApiMessageAndVersion record = new ApiMessageAndVersion(new KVRecord() + .setKeyValues(Collections.singletonList(new KVRecord.KeyValue() + .setKey(FAILOVER_KEY) + .setValue(encoded.getBytes(StandardCharsets.UTF_8)))), (short) 0); + return ControllerResult.of(Collections.singletonList(record), null); + } + + public void replay(KVRecord record) { + if (!record.keyValues().isEmpty()) { + KVRecord.KeyValue kv = record.keyValues().get(0); + if (!FAILOVER_KEY.equals(kv.key())) { + return; + } + FailoverContext[] contexts = JsonUtils.decode(new String(kv.value(), StandardCharsets.UTF_8), FailoverContext[].class); + Map failoverContexts = new HashMap<>(contexts.length); + for (FailoverContext ctx : contexts) { + failoverContexts.put(ctx.getFailedNode(), ctx); + } + this.failoverContexts.set(failoverContexts); + } + } + + private List allNodes() { + // source of truth for node epoch + Map nodeEpochMap = streamControlManager.nodesMetadata().values().stream() + .collect(Collectors.toMap( + org.apache.kafka.controller.stream.NodeRuntimeMetadata::getNodeId, + org.apache.kafka.controller.stream.NodeRuntimeMetadata::getNodeEpoch + )); + + return nodeControlManager.getMetadata().stream() + // Generally, any node in nodeControlManager should have a corresponding node in streamControlManager. + // However, there is a very short period of time when a node is in nodeControlManager but not in streamControlManager when the node first joins the cluster. + // In this case, we just ignore the node. + .filter(node -> nodeEpochMap.containsKey(node.getNodeId())) + .map(node -> new NodeRuntimeMetadata( + node.getNodeId(), + // There are node epochs in both streamControlManager and nodeControlManager, and they are the same in most cases. + // However, in some rare cases, the node epoch in streamControlManager may be updated earlier than the node epoch in nodeControlManager. + // So we use the node epoch in streamControlManager as the source of truth. + nodeEpochMap.get(node.getNodeId()), + node.getWalConfig(), + node.getTags(), + nodeControlManager.state(node.getNodeId()), + nodeControlManager.hasOpeningStreams(node.getNodeId()) + )) + .collect(Collectors.toList()); + } + + private List addNewFailover(Map failoverContexts, + List failedWalList) { + List newFailover = new LinkedList<>(); + List brokerIds = null; + Map failoverCounts = new HashMap<>(); + // round-robin assign new failover + int assignIndex = 0; + for (FailedWal failedWal : failedWalList) { + if (failoverContexts.containsKey(failedWal.node())) { + continue; + } + if (brokerIds == null) { + // lazy init + brokerIds = clusterControlManager.getActiveBrokers().stream().map(BrokerRegistration::id).collect(Collectors.toList()); + failoverContexts.forEach((n, ctx) -> failoverCounts.merge(ctx.getTarget(), 1, Integer::sum)); + } + boolean found = false; + for (int i = 0; i < brokerIds.size(); i++, assignIndex++) { + int brokerId = brokerIds.get(assignIndex % brokerIds.size()); + + if (brokerId == failedWal.nodeId()) { + // skip the failed node itself + continue; + } + + int attachedCount = Optional.ofNullable(failoverCounts.get(brokerId)).orElse(0); + if (attachedCount >= MAX_FAILOVER_COUNT_IN_TARGET_NODE) { + continue; + } + failoverCounts.merge(brokerId, 1, Integer::sum); + newFailover.add(failedWal.toFailoverContext(brokerId)); + found = true; + break; + } + if (!found) { + LOGGER.warn("No broker available for failover, failedWal={}", failedWal); + } + } + return newFailover; + } + + private List delCompleted(Map failoverContexts, + List failedWalList) { + Set failedNodeSet = failedWalList.stream() + .map(FailedWal::node) + .collect(Collectors.toSet()); + return failoverContexts.keySet().stream() + .filter(node -> !failedNodeSet.contains(node)) + .collect(Collectors.toList()); + } + + /** + * Reset the failover context whose target node is not alive + */ + private List resetNotAlive(Map failoverContexts) { + List reset = new LinkedList<>(); + for (FailoverContext ctx : failoverContexts.values()) { + if (!clusterControlManager.isActive(ctx.getTarget())) { + reset.add(ctx.getFailedNode()); + } + } + return reset; + } + + @Override + public void close() throws Exception { + scheduler.shutdown(); + } + + private static List getFailedWal(List allNodes) { + List result = allNodes.stream() + .filter(NodeRuntimeMetadata::shouldFailover) + .map(DefaultFailedWal::from) + .collect(Collectors.toCollection(ArrayList::new)); + return result; + } +} diff --git a/core/src/main/java/kafka/automq/failover/FailoverListener.java b/core/src/main/java/kafka/automq/failover/FailoverListener.java new file mode 100644 index 0000000000..1d482952a9 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/FailoverListener.java @@ -0,0 +1,149 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +import kafka.automq.utils.JsonUtils; + +import org.apache.kafka.image.KVDelta; +import org.apache.kafka.image.MetadataDelta; +import org.apache.kafka.image.MetadataImage; +import org.apache.kafka.image.loader.LoaderManifest; +import org.apache.kafka.image.publisher.MetadataPublisher; + +import com.automq.stream.api.Client; +import com.automq.stream.s3.failover.FailoverRequest; +import com.automq.stream.utils.ThreadUtils; +import com.automq.stream.utils.Threads; +import com.google.common.collect.Sets; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +public class FailoverListener implements MetadataPublisher, AutoCloseable { + private static final Logger LOGGER = LoggerFactory.getLogger(FailoverListener.class); + private final int nodeId; + private final Map recovering = new ConcurrentHashMap<>(); + + private final Client client; + + private final ScheduledExecutorService scheduler = Threads.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("failover-listener-%d", true), LOGGER); + + public FailoverListener(int nodeId, Client client) { + this.nodeId = nodeId; + this.client = client; + } + + @Override + public void onMetadataUpdate(MetadataDelta delta, MetadataImage image, LoaderManifest manifest) { + try { + getContexts(delta).ifPresent(this::onContextsChange); + } catch (Throwable e) { + LOGGER.error("failover listener fail", e); + } + } + + /** + * Get the failover contexts from the given metadata delta. + * It returns empty if there is no update for the failover key. + */ + private Optional getContexts(MetadataDelta delta) { + return Optional.ofNullable(delta.kvDelta()) + .map(KVDelta::changedKV) + .map(kv -> kv.get(FailoverConstants.FAILOVER_KEY)) + .map(this::decodeContexts); + } + + private FailoverContext[] decodeContexts(ByteBuffer byteBuffer) { + ByteBuffer slice = byteBuffer.slice(); + byte[] data = new byte[slice.remaining()]; + slice.get(data); + return JsonUtils.decode(new String(data, StandardCharsets.UTF_8), FailoverContext[].class); + } + + private void onContextsChange(FailoverContext[] contexts) { + Set oldFailedNodes = recovering.keySet(); + Set newFailedNodes = Arrays.stream(contexts) + .filter(ctx -> ctx.getTarget() == nodeId) + .map(FailoverContext::getFailedNode) + .collect(Collectors.toSet()); + + Set completedNodes = Sets.difference(oldFailedNodes, newFailedNodes); + completedNodes.forEach(recovering::remove); + + Set needFailoverNodes = Sets.difference(newFailedNodes, oldFailedNodes); + for (FailoverContext context : contexts) { + FailedNode failedNode = context.getFailedNode(); + if (needFailoverNodes.contains(failedNode)) { + recovering.put(failedNode, context); + failover(context); + } + } + } + + private void failover(FailoverContext context) { + scheduler.execute(() -> failover0(context)); + } + + private void failover0(FailoverContext context) { + failover0(context, 0); + } + + + private void failover0(FailoverContext context, int retryCount) { + try { + if (!recovering.containsKey(context.getFailedNode())) { + return; + } + LOGGER.info("[FAILOVER] start with context={}, retryCount={}", context, retryCount); + FailoverRequest request = new FailoverRequest(); + request.setNodeId(context.getNodeId()); + request.setNodeEpoch(context.getNodeEpoch()); + request.setKraftWalConfigs(context.getKraftWalConfigs()); + client.failover(request).get(); + LOGGER.info("[FAILOVER] complete with context={}, retryCount={}", context, retryCount); + } catch (Throwable e) { + int retryDelay = Math.min(1 << retryCount, 60); + LOGGER.warn("[FAILOVER] fail, retry later. context={}, retryCount={}, retryDelay={}s", context, retryCount, retryDelay, e); + if (recovering.containsKey(context.getFailedNode())) { + scheduler.schedule(() -> failover0(context, retryCount + 1), retryDelay, TimeUnit.SECONDS); + } + } + } + + @Override + public String name() { + return "failover-listener"; + } + + @Override + public void close() throws Exception { + } +} diff --git a/core/src/main/java/kafka/automq/failover/NodeRuntimeMetadata.java b/core/src/main/java/kafka/automq/failover/NodeRuntimeMetadata.java new file mode 100644 index 0000000000..7e28b66e05 --- /dev/null +++ b/core/src/main/java/kafka/automq/failover/NodeRuntimeMetadata.java @@ -0,0 +1,151 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +import org.apache.kafka.controller.ClusterControlManager; +import org.apache.kafka.controller.stream.NodeMetadata; +import org.apache.kafka.controller.stream.NodeState; + +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * NodeRuntimeMetadata is a runtime view of a node's metadata. + * + * @see NodeMetadata + */ +public final class NodeRuntimeMetadata { + + /** + * Expect the Node ID of a broker (as opposed to a controller) will start from 1000. + * + * @see ClusterControlManager#getNextNodeId() + */ + private static final int MAX_CONTROLLER_ID = 1000 - 1; + private static final long DONT_FAILOVER_AFTER_NEW_EPOCH_MS = TimeUnit.MINUTES.toMillis(1); + private final int id; + private final long epoch; + private final String walConfigs; + private final Map tags; + private final NodeState state; + private final boolean hasOpeningStreams; + + /** + * + */ + public NodeRuntimeMetadata(int id, long epoch, String walConfigs, Map tags, NodeState state, + boolean hasOpeningStreams) { + this.id = id; + this.epoch = epoch; + this.walConfigs = walConfigs; + this.tags = tags; + this.state = state; + this.hasOpeningStreams = hasOpeningStreams; + } + + public boolean shouldFailover() { + return isFenced() && hasOpeningStreams + // The node epoch is the start timestamp of node. + // We need to avoid failover just after node restart. + // The node may take some time to recover its data. + && System.currentTimeMillis() - epoch > DONT_FAILOVER_AFTER_NEW_EPOCH_MS; + } + + public boolean isFenced() { + return NodeState.FENCED == state; + } + + public boolean isActive() { + return NodeState.ACTIVE == state; + } + + public boolean isController() { + return id <= MAX_CONTROLLER_ID; + } + + public static boolean isController(int nodeId) { + return nodeId <= MAX_CONTROLLER_ID; + } + + private String getTagOrThrow(String key) { + String value = tags.get(key); + if (value == null) { + throw new IllegalStateException(String.format("Node %d is missing tag %s, tags: %s", id, key, tags)); + } + return value; + } + + public int id() { + return id; + } + + public long epoch() { + return epoch; + } + + public String walConfigs() { + return walConfigs; + } + + public Map tags() { + return tags; + } + + public NodeState state() { + return state; + } + + public boolean hasOpeningStreams() { + return hasOpeningStreams; + } + + @Override + public boolean equals(Object obj) { + if (obj == this) + return true; + if (obj == null || obj.getClass() != this.getClass()) + return false; + var that = (NodeRuntimeMetadata) obj; + return this.id == that.id && + this.epoch == that.epoch && + Objects.equals(this.walConfigs, that.walConfigs) && + Objects.equals(this.tags, that.tags) && + Objects.equals(this.state, that.state) && + this.hasOpeningStreams == that.hasOpeningStreams; + } + + @Override + public int hashCode() { + return Objects.hash(id, epoch, walConfigs, tags, state, hasOpeningStreams); + } + + @Override + public String toString() { + return "NodeRuntimeMetadata[" + + "id=" + id + ", " + + "epoch=" + epoch + ", " + + "walConfigs=" + walConfigs + ", " + + "tags=" + tags + ", " + + "state=" + state + ", " + + "hasOpeningStreams=" + hasOpeningStreams + ']'; + } + +} diff --git a/core/src/main/java/kafka/automq/interceptor/ClientIdKey.java b/core/src/main/java/kafka/automq/interceptor/ClientIdKey.java new file mode 100644 index 0000000000..8be2e09728 --- /dev/null +++ b/core/src/main/java/kafka/automq/interceptor/ClientIdKey.java @@ -0,0 +1,25 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.interceptor; + +public class ClientIdKey { + public static final String AVAILABILITY_ZONE = "automq_az"; + public static final String CLIENT_TYPE = "automq_type"; +} diff --git a/core/src/main/java/kafka/automq/interceptor/ClientIdMetadata.java b/core/src/main/java/kafka/automq/interceptor/ClientIdMetadata.java new file mode 100644 index 0000000000..53d0f535ec --- /dev/null +++ b/core/src/main/java/kafka/automq/interceptor/ClientIdMetadata.java @@ -0,0 +1,96 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.interceptor; + +import com.automq.stream.utils.URIUtils; + +import java.net.InetAddress; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public class ClientIdMetadata { + private final String clientId; + private final Map> metadata; + private final InetAddress clientAddress; + private final String connectionId; + + private ClientIdMetadata(String clientId, InetAddress clientAddress, String connectionId) { + this.clientId = clientId; + this.metadata = URIUtils.splitQuery(clientId); + this.clientAddress = clientAddress; + this.connectionId = connectionId; + } + + public static ClientIdMetadata of(String clientId) { + return new ClientIdMetadata(clientId, null, null); + } + + public static ClientIdMetadata of(String clientId, InetAddress clientAddress, String connectionId) { + return new ClientIdMetadata(clientId, clientAddress, connectionId); + } + + public String rack() { + List list = metadata.get(ClientIdKey.AVAILABILITY_ZONE); + if (list == null || list.isEmpty()) { + return null; + } + return list.get(0); + } + + public ClientType clientType() { + List list = metadata.get(ClientIdKey.CLIENT_TYPE); + if (list == null || list.isEmpty()) { + return null; + } + return ClientType.parse(list.get(0)); + } + + public String clientId() { + return clientId; + } + + public String connectionId() { + return connectionId; + } + + public InetAddress clientAddress() { + return clientAddress; + } + + public List metadata(String key) { + List value = metadata.get(key); + return Objects.requireNonNullElse(value, Collections.emptyList()); + } + + public void metadata(String key, List valueList) { + metadata.put(key, valueList); + } + + @Override + public String toString() { + if (clientAddress == null) { + return clientId; + } else { + return clientId + "/" + clientAddress.getHostAddress(); + } + } +} diff --git a/core/src/main/java/kafka/automq/interceptor/ClientType.java b/core/src/main/java/kafka/automq/interceptor/ClientType.java new file mode 100644 index 0000000000..b305b02caf --- /dev/null +++ b/core/src/main/java/kafka/automq/interceptor/ClientType.java @@ -0,0 +1,36 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.interceptor; + +public enum ClientType { + PRODUCER, CONSUMER; + + public static ClientType parse(String str) { + switch (str) { + case "producer": + return PRODUCER; + case "consumer": + return CONSUMER; + default: + return null; + } + } + +} diff --git a/core/src/main/java/kafka/automq/interceptor/NoopTrafficInterceptor.java b/core/src/main/java/kafka/automq/interceptor/NoopTrafficInterceptor.java new file mode 100644 index 0000000000..3587048e46 --- /dev/null +++ b/core/src/main/java/kafka/automq/interceptor/NoopTrafficInterceptor.java @@ -0,0 +1,76 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.interceptor; + +import kafka.server.MetadataCache; +import kafka.server.streamaspect.ElasticKafkaApis; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.message.AutomqZoneRouterRequestData; +import org.apache.kafka.common.message.MetadataResponseData; +import org.apache.kafka.common.network.ListenerName; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; + +import com.automq.stream.utils.FutureUtil; + +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +public class NoopTrafficInterceptor implements TrafficInterceptor { + private final ElasticKafkaApis kafkaApis; + private final MetadataCache metadataCache; + + public NoopTrafficInterceptor(ElasticKafkaApis kafkaApis, MetadataCache metadataCache) { + this.kafkaApis = kafkaApis; + this.metadataCache = metadataCache; + } + + @Override + public void close() { + + } + + @Override + public void handleProduceRequest(ProduceRequestArgs args) { + kafkaApis.handleProduceAppendJavaCompatible(args); + } + + @Override + public CompletableFuture handleZoneRouterRequest(AutomqZoneRouterRequestData request) { + return FutureUtil.failedFuture(new UnsupportedOperationException()); + } + + @Override + public List handleMetadataResponse(ClientIdMetadata clientId, + List topics) { + return topics; + } + + @Override + public Optional getLeaderNode(int leaderId, ClientIdMetadata clientId, String listenerName) { + scala.Option opt = metadataCache.getAliveBrokerNode(leaderId, new ListenerName(listenerName)); + if (opt.isEmpty()) { + return Optional.empty(); + } else { + return Optional.of(opt.get()); + } + } +} diff --git a/core/src/main/java/kafka/automq/interceptor/ProduceRequestArgs.java b/core/src/main/java/kafka/automq/interceptor/ProduceRequestArgs.java new file mode 100644 index 0000000000..0273c7ce49 --- /dev/null +++ b/core/src/main/java/kafka/automq/interceptor/ProduceRequestArgs.java @@ -0,0 +1,184 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.interceptor; + +import kafka.server.RequestLocal; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.RecordValidationStats; +import org.apache.kafka.common.requests.ProduceResponse; + +import java.util.Map; +import java.util.function.Consumer; + +public class ProduceRequestArgs { + private final short apiVersion; + private final ClientIdMetadata clientId; + private final int timeout; + private final short requiredAcks; + private final boolean internalTopicsAllowed; + private final String transactionId; + private final Map entriesPerPartition; + private final Consumer> responseCallback; + private final Consumer> recordValidationStatsCallback; + private final RequestLocal requestLocal; + + public ProduceRequestArgs(short apiVersion, ClientIdMetadata id, int timeout, short requiredAcks, boolean allowed, + String transactionId, Map partition, + Consumer> callback, + Consumer> statsCallback, RequestLocal local) { + this.apiVersion = apiVersion; + this.clientId = id; + this.timeout = timeout; + this.requiredAcks = requiredAcks; + this.internalTopicsAllowed = allowed; + this.transactionId = transactionId; + this.entriesPerPartition = partition; + this.responseCallback = callback; + this.recordValidationStatsCallback = statsCallback; + this.requestLocal = local; + } + + public short apiVersion() { + return apiVersion; + } + + public ClientIdMetadata clientId() { + return clientId; + } + + public int timeout() { + return timeout; + } + + public short requiredAcks() { + return requiredAcks; + } + + public boolean internalTopicsAllowed() { + return internalTopicsAllowed; + } + + public String transactionId() { + return transactionId; + } + + public Map entriesPerPartition() { + return entriesPerPartition; + } + + public Consumer> responseCallback() { + return responseCallback; + } + + public Consumer> recordValidationStatsCallback() { + return recordValidationStatsCallback; + } + + public RequestLocal requestLocal() { + return requestLocal; + } + + public Builder toBuilder() { + return new Builder() + .apiVersion(apiVersion) + .clientId(clientId) + .timeout(timeout) + .requiredAcks(requiredAcks) + .internalTopicsAllowed(internalTopicsAllowed) + .transactionId(transactionId) + .entriesPerPartition(entriesPerPartition) + .responseCallback(responseCallback) + .recordValidationStatsCallback(recordValidationStatsCallback) + .requestLocal(requestLocal); + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private short apiVersion; + private ClientIdMetadata clientId; + private int timeout; + private short requiredAcks; + private boolean internalTopicsAllowed; + private String transactionId; + private Map entriesPerPartition; + private Consumer> responseCallback; + private Consumer> recordValidationStatsCallback; + private RequestLocal requestLocal; + + public Builder apiVersion(short apiVersion) { + this.apiVersion = apiVersion; + return this; + } + + public Builder clientId(ClientIdMetadata clientId) { + this.clientId = clientId; + return this; + } + + public Builder timeout(int timeout) { + this.timeout = timeout; + return this; + } + + public Builder requiredAcks(short requiredAcks) { + this.requiredAcks = requiredAcks; + return this; + } + + public Builder internalTopicsAllowed(boolean internalTopicsAllowed) { + this.internalTopicsAllowed = internalTopicsAllowed; + return this; + } + + public Builder transactionId(String transactionId) { + this.transactionId = transactionId; + return this; + } + + public Builder entriesPerPartition(Map entriesPerPartition) { + this.entriesPerPartition = entriesPerPartition; + return this; + } + + public Builder responseCallback(Consumer> responseCallback) { + this.responseCallback = responseCallback; + return this; + } + + public Builder recordValidationStatsCallback(Consumer> recordValidationStatsCallback) { + this.recordValidationStatsCallback = recordValidationStatsCallback; + return this; + } + + public Builder requestLocal(RequestLocal requestLocal) { + this.requestLocal = requestLocal; + return this; + } + + public ProduceRequestArgs build() { + return new ProduceRequestArgs(apiVersion, clientId, timeout, requiredAcks, internalTopicsAllowed, transactionId, entriesPerPartition, responseCallback, recordValidationStatsCallback, requestLocal); + } + } +} diff --git a/core/src/main/java/kafka/automq/interceptor/TrafficInterceptor.java b/core/src/main/java/kafka/automq/interceptor/TrafficInterceptor.java new file mode 100644 index 0000000000..12f27839e1 --- /dev/null +++ b/core/src/main/java/kafka/automq/interceptor/TrafficInterceptor.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.interceptor; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.message.AutomqZoneRouterRequestData; +import org.apache.kafka.common.message.MetadataResponseData; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; + +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +public interface TrafficInterceptor { + + void close(); + + void handleProduceRequest(ProduceRequestArgs args); + + CompletableFuture handleZoneRouterRequest(AutomqZoneRouterRequestData request); + + List handleMetadataResponse(ClientIdMetadata clientId, + List topics); + + Optional getLeaderNode(int leaderId, ClientIdMetadata clientId, String listenerName); + +} diff --git a/core/src/main/java/kafka/automq/kafkalinking/KafkaLinkingManager.java b/core/src/main/java/kafka/automq/kafkalinking/KafkaLinkingManager.java new file mode 100644 index 0000000000..022b52db33 --- /dev/null +++ b/core/src/main/java/kafka/automq/kafkalinking/KafkaLinkingManager.java @@ -0,0 +1,32 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.kafkalinking; + +import org.apache.kafka.common.TopicPartition; + +import java.util.Set; + +public interface KafkaLinkingManager { + void addPartitions(Set partitions); + + void removePartitions(Set partitions); + + void shutdown(); +} diff --git a/core/src/main/java/kafka/automq/partition/snapshot/ConfirmWalDataDelta.java b/core/src/main/java/kafka/automq/partition/snapshot/ConfirmWalDataDelta.java new file mode 100644 index 0000000000..eb0f442df7 --- /dev/null +++ b/core/src/main/java/kafka/automq/partition/snapshot/ConfirmWalDataDelta.java @@ -0,0 +1,191 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.partition.snapshot; + +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; + +import com.automq.stream.s3.ConfirmWAL; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.RecordOffset; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicInteger; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +/** + * Maintains a bounded, in-memory delta of recent WAL appends so snapshot responses can + * piggy-back fresh data instead of forcing clients to replay the physical WAL. + * + *

Responsibilities + *

    + *
  • Subscribe to {@link ConfirmWAL} append events and retain the encoded + * {@link StreamRecordBatch} payloads while they are eligible for delta export.
    • + *
  • Track confirm offsets and expose them via {@link #handle(short, AutomqGetPartitionSnapshotResponseData)}.
    • + *
  • Serialize buffered batches into {@code confirmWalDeltaData} for request versions + * >= 2, or signal that callers must replay the WAL otherwise.
    • + *
  • Enforce {@link #MAX_RECORDS_BUFFER_SIZE} so the delta cache remains lightweight.
    • + *
+ * + *

State machine + *

    + *
  • {@link #STATE_NOT_SYNC}: Buffer content is discarded (e.g. overflow) and only confirm + * offsets are returned until new appends arrive.
    • + *
  • {@link #STATE_SYNCING}: Buffered records are eligible to be drained and turned into a + * delta payload when {@link #handle(short, AutomqGetPartitionSnapshotResponseData)} runs.
    • + *
  • {@link #STATE_CLOSED}: Listener is torn down and ignores subsequent appends.
    • + *
+ * + *

Concurrency and lifecycle + *

    + *
  • All public methods are synchronized to guard the state machine, queue, and + * {@link #lastConfirmOffset} tracking.
    • + *
  • Buffered batches are reference-counted; ownership transfers to this class until the + * delta is emitted or the buffer is dropped/closed.
    • + *
  • {@link #close()} must be invoked when the owning {@link PartitionSnapshotsManager.Session} ends to release buffers + * and remove the {@link ConfirmWAL.AppendListener}.
    • + *
+ * + *

Snapshot interaction + *

    + *
  • {@link #handle(short, AutomqGetPartitionSnapshotResponseData)} always updates + * {@code confirmWalEndOffset} and, when possible, attaches {@code confirmWalDeltaData}.
    • + *
  • A {@code null} delta signals the client must replay the WAL, whereas an empty byte array + * indicates no new data but confirms offsets.
    • + *
  • When the aggregated encoded bytes would exceed {@link #MAX_RECORDS_BUFFER_SIZE}, the + * buffer is dropped and state resets to {@link #STATE_NOT_SYNC}.
    • + *
+ */ +public class ConfirmWalDataDelta implements ConfirmWAL.AppendListener { + static final int STATE_NOT_SYNC = 0; + static final int STATE_SYNCING = 1; + static final int STATE_CLOSED = 9; + static final int MAX_RECORDS_BUFFER_SIZE = 32 * 1024; // 32KiB + private final ConfirmWAL confirmWAL; + + private final ConfirmWAL.ListenerHandle listenerHandle; + final BlockingQueue records = new LinkedBlockingQueue<>(); + final AtomicInteger size = new AtomicInteger(0); + + private RecordOffset lastConfirmOffset = null; + + int state = STATE_NOT_SYNC; + + public ConfirmWalDataDelta(ConfirmWAL confirmWAL) { + this.confirmWAL = confirmWAL; + this.listenerHandle = confirmWAL.addAppendListener(this); + } + + public synchronized void close() { + this.state = STATE_CLOSED; + this.listenerHandle.close(); + records.forEach(r -> r.record.release()); + records.clear(); + } + + public void handle(short requestVersion, + AutomqGetPartitionSnapshotResponseData resp) { + RecordOffset newConfirmOffset = null; + List delta = null; + synchronized (this) { + if (state == STATE_NOT_SYNC) { + List drainedRecords = new ArrayList<>(records.size()); + records.drainTo(drainedRecords); + size.addAndGet(-drainedRecords.stream().mapToInt(r -> r.record.encoded().readableBytes()).sum()); + if (!drainedRecords.isEmpty()) { + RecordOffset deltaConfirmOffset = drainedRecords.get(drainedRecords.size() - 1).nextOffset(); + if (lastConfirmOffset == null || deltaConfirmOffset.compareTo(lastConfirmOffset) > 0) { + newConfirmOffset = deltaConfirmOffset; + state = STATE_SYNCING; + } + drainedRecords.forEach(r -> r.record.release()); + } + } else if (state == STATE_SYNCING) { + delta = new ArrayList<>(records.size()); + + records.drainTo(delta); + size.addAndGet(-delta.stream().mapToInt(r -> r.record.encoded().readableBytes()).sum()); + newConfirmOffset = delta.isEmpty() ? lastConfirmOffset : delta.get(delta.size() - 1).nextOffset(); + } + if (newConfirmOffset == null) { + newConfirmOffset = confirmWAL.confirmOffset(); + } + this.lastConfirmOffset = newConfirmOffset; + } + resp.setConfirmWalEndOffset(newConfirmOffset.bufferAsBytes()); + if (delta != null) { + int size = delta.stream().mapToInt(r -> r.record.encoded().readableBytes()).sum(); + byte[] data = new byte[size]; + ByteBuf buf = Unpooled.wrappedBuffer(data).clear(); + delta.forEach(r -> { + buf.writeBytes(r.record.encoded()); + r.record.release(); + }); + if (requestVersion >= 2) { + // The confirmWalDeltaData is only supported in request version >= 2 + resp.setConfirmWalDeltaData(data); + } + } else { + if (requestVersion >= 2) { + // - Null means the client needs replay from the physical WAL + // - Empty means there is no delta data. + resp.setConfirmWalDeltaData(null); + } + } + } + + @Override + public synchronized void onAppend(StreamRecordBatch record, RecordOffset recordOffset, + RecordOffset nextOffset) { + if (state == STATE_CLOSED) { + return; + } + record.retain(); + records.add(new RecordExt(record, recordOffset, nextOffset)); + if (size.addAndGet(record.encoded().readableBytes()) > MAX_RECORDS_BUFFER_SIZE) { + // If the buffer is full, drop all records and switch to NOT_SYNC state. + // It's cheaper to replay from the physical WAL instead of transferring the data by network. + state = STATE_NOT_SYNC; + records.forEach(r -> r.record.release()); + records.clear(); + size.set(0); + } + } + + record RecordExt(StreamRecordBatch record, RecordOffset recordOffset, RecordOffset nextOffset) { + } + + public static List decodeDeltaRecords(byte[] data) { + if (data == null) { + return null; + } + List records = new ArrayList<>(); + ByteBuf buf = Unpooled.wrappedBuffer(data); + while (buf.readableBytes() > 0) { + StreamRecordBatch record = StreamRecordBatch.parse(buf, false); + records.add(record); + } + return records; + } +} diff --git a/core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotVersion.java b/core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotVersion.java new file mode 100644 index 0000000000..d44f978a8b --- /dev/null +++ b/core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotVersion.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.partition.snapshot; + +import java.util.Objects; + +public class PartitionSnapshotVersion { + // increment every time when there is a segment change + private int segmentsVersion; + // increment every time when there is a new record + private int recordsVersion; + + private PartitionSnapshotVersion(int segmentsVersion, int recordsVersion) { + this.segmentsVersion = segmentsVersion; + this.recordsVersion = recordsVersion; + } + + public static PartitionSnapshotVersion create() { + return new PartitionSnapshotVersion(0, 0); + } + + public int segmentsVersion() { + return segmentsVersion; + } + + public int recordsVersion() { + return recordsVersion; + } + + public PartitionSnapshotVersion incrementSegmentsVersion() { + segmentsVersion++; + return this; + } + + public PartitionSnapshotVersion incrementRecordsVersion() { + recordsVersion++; + return this; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + PartitionSnapshotVersion version = (PartitionSnapshotVersion) o; + return segmentsVersion == version.segmentsVersion && recordsVersion == version.recordsVersion; + } + + @Override + public int hashCode() { + return Objects.hash(segmentsVersion, recordsVersion); + } + + public PartitionSnapshotVersion copy() { + return new PartitionSnapshotVersion(segmentsVersion, recordsVersion); + } + +} diff --git a/core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotsManager.java b/core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotsManager.java new file mode 100644 index 0000000000..4990598d0d --- /dev/null +++ b/core/src/main/java/kafka/automq/partition/snapshot/PartitionSnapshotsManager.java @@ -0,0 +1,378 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.partition.snapshot; + +import kafka.automq.AutoMQConfig; +import kafka.cluster.LogEventListener; +import kafka.cluster.Partition; +import kafka.cluster.PartitionListener; +import kafka.log.streamaspect.ElasticLogMeta; +import kafka.log.streamaspect.ElasticStreamSegmentMeta; +import kafka.log.streamaspect.SliceRange; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotRequestData; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData.LogMetadata; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData.PartitionSnapshot; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData.SegmentMetadata; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData.StreamMappingCollection; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData.Topic; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData.TopicCollection; +import org.apache.kafka.common.requests.s3.AutomqGetPartitionSnapshotRequest; +import org.apache.kafka.common.requests.s3.AutomqGetPartitionSnapshotResponse; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.server.common.automq.AutoMQVersion; +import org.apache.kafka.storage.internals.log.LogOffsetMetadata; +import org.apache.kafka.storage.internals.log.TimestampOffset; + +import com.automq.stream.s3.ConfirmWAL; +import com.automq.stream.utils.Threads; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; +import java.util.stream.Collectors; + +import io.netty.util.concurrent.FastThreadLocal; + +public class PartitionSnapshotsManager { + private static final int NOOP_SESSION_ID = 0; + private final Map sessions = new HashMap<>(); + private final List snapshotVersions = new CopyOnWriteArrayList<>(); + private final Time time; + private final ConfirmWAL confirmWAL; + + public PartitionSnapshotsManager(Time time, AutoMQConfig config, ConfirmWAL confirmWAL, + Supplier versionGetter) { + this.time = time; + this.confirmWAL = confirmWAL; + if (config.zoneRouterChannels().isPresent()) { + Threads.COMMON_SCHEDULER.scheduleWithFixedDelay(this::cleanExpiredSessions, 1, 1, TimeUnit.MINUTES); + Threads.COMMON_SCHEDULER.scheduleWithFixedDelay(() -> { + // In ZERO_ZONE_V0 we need to fast commit the WAL data to KRaft, + // then another nodes could replay the SSO to support snapshot read. + if (!versionGetter.get().isZeroZoneV2Supported()) { + confirmWAL.commit(0, false); + } + }, 1, 1, TimeUnit.SECONDS); + } + } + + public void onPartitionOpen(Partition partition) { + PartitionWithVersion partitionWithVersion = new PartitionWithVersion(partition, PartitionSnapshotVersion.create()); + snapshotVersions.add(partitionWithVersion); + partition.maybeAddListener(newPartitionListener(partitionWithVersion)); + partition.addLogEventListener(newLogEventListener(partitionWithVersion)); + } + + public void onPartitionClose(Partition partition) { + snapshotVersions.removeIf(p -> p.partition == partition); + synchronized (this) { + sessions.values().forEach(s -> s.onPartitionClose(partition)); + } + } + + public CompletableFuture handle(AutomqGetPartitionSnapshotRequest request) { + Session session; + boolean newSession = false; + synchronized (this) { + AutomqGetPartitionSnapshotRequestData requestData = request.data(); + int sessionId = requestData.sessionId(); + int sessionEpoch = requestData.sessionEpoch(); + session = sessions.get(sessionId); + if (sessionId == NOOP_SESSION_ID + || session == null + || (sessionEpoch != session.sessionEpoch())) { + if (session != null) { + sessions.remove(sessionId); + } + sessionId = nextSessionId(); + session = new Session(sessionId); + sessions.put(sessionId, session); + newSession = true; + } + } + return session.snapshotsDelta(request, request.data().requestCommit() || newSession); + } + + private synchronized int nextSessionId() { + int id; + do { + id = ThreadLocalRandom.current().nextInt(1, Integer.MAX_VALUE); + } + while (sessions.containsKey(id) || id == NOOP_SESSION_ID); + return id; + } + + private synchronized void cleanExpiredSessions() { + sessions.values().removeIf(s -> { + boolean expired = s.expired(); + if (expired) { + s.close(); + } + return expired; + }); + } + + class Session { + private static final short ZERO_ZONE_V0_REQUEST_VERSION = (short) 0; + private static final FastThreadLocal>> COMPLETE_CF_LIST_LOCAL = new FastThreadLocal<>() { + @Override + protected List> initialValue() { + return new ArrayList<>(); + } + }; + private final int sessionId; + private int sessionEpoch = 0; + private final Map synced = new HashMap<>(); + private final List removed = new ArrayList<>(); + private long lastGetSnapshotsTimestamp = time.milliseconds(); + private final Set> inflightCommitCfSet = ConcurrentHashMap.newKeySet(); + private final ConfirmWalDataDelta delta; + + public Session(int sessionId) { + this.sessionId = sessionId; + this.delta = new ConfirmWalDataDelta(confirmWAL); + } + + public synchronized void close() { + delta.close(); + } + + public synchronized int sessionEpoch() { + return sessionEpoch; + } + + public synchronized CompletableFuture snapshotsDelta( + AutomqGetPartitionSnapshotRequest request, boolean requestCommit) { + AutomqGetPartitionSnapshotResponseData resp = new AutomqGetPartitionSnapshotResponseData(); + sessionEpoch++; + lastGetSnapshotsTimestamp = time.milliseconds(); + resp.setSessionId(sessionId); + resp.setSessionEpoch(sessionEpoch); + long finalSessionEpoch = sessionEpoch; + CompletableFuture collectPartitionSnapshotsCf; + if (!requestCommit && inflightCommitCfSet.isEmpty()) { + collectPartitionSnapshotsCf = collectPartitionSnapshots(request.data().version(), resp); + } else { + collectPartitionSnapshotsCf = CompletableFuture.completedFuture(null); + } + boolean newSession = finalSessionEpoch == 1; + return collectPartitionSnapshotsCf + .thenApply(nil -> { + if (request.data().version() > ZERO_ZONE_V0_REQUEST_VERSION) { + if (newSession) { + // return the WAL config in the session first response + resp.setConfirmWalConfig(confirmWAL.uri()); + } + delta.handle(request.version(), resp); + } + if (requestCommit) { + // Commit after generating the snapshots. + // Then the snapshot-read partitions could read from snapshot-read cache or block cache. + CompletableFuture commitCf = newSession ? + // The proxy node's first snapshot-read request needs to commit immediately to ensure the data could be read. + confirmWAL.commit(0, false) + // The proxy node's snapshot-read cache isn't enough to hold the 'uncommitted' data, + // so the proxy node request a commit to ensure the data could be read from block cache. + : confirmWAL.commit(1000, false); + inflightCommitCfSet.add(commitCf); + commitCf.whenComplete((rst, ex) -> inflightCommitCfSet.remove(commitCf)); + } + return new AutomqGetPartitionSnapshotResponse(resp); + }); + } + + public synchronized void onPartitionClose(Partition partition) { + removed.add(partition); + } + + public synchronized boolean expired() { + return time.milliseconds() - lastGetSnapshotsTimestamp > 60000; + } + + private CompletableFuture collectPartitionSnapshots(short funcVersion, + AutomqGetPartitionSnapshotResponseData resp) { + Map> topic2partitions = new HashMap<>(); + List> completeCfList = COMPLETE_CF_LIST_LOCAL.get(); + completeCfList.clear(); + removed.forEach(partition -> { + PartitionSnapshotVersion version = synced.remove(partition); + if (version != null) { + List partitionSnapshots = topic2partitions.computeIfAbsent(partition.topicId().get(), topic -> new ArrayList<>()); + partitionSnapshots.add(snapshot(funcVersion, partition, version, null, completeCfList)); + } + }); + removed.clear(); + + snapshotVersions.forEach(p -> { + PartitionSnapshotVersion oldVersion = synced.get(p.partition); + if (!Objects.equals(p.version, oldVersion)) { + List partitionSnapshots = topic2partitions.computeIfAbsent(p.partition.topicId().get(), topic -> new ArrayList<>()); + PartitionSnapshotVersion newVersion = p.version.copy(); + PartitionSnapshot partitionSnapshot = snapshot(funcVersion, p.partition, oldVersion, newVersion, completeCfList); + partitionSnapshots.add(partitionSnapshot); + synced.put(p.partition, newVersion); + } + }); + TopicCollection topics = new TopicCollection(); + topic2partitions.forEach((topicId, partitions) -> { + Topic topic = new Topic(); + topic.setTopicId(topicId); + topic.setPartitions(partitions); + topics.add(topic); + }); + resp.setTopics(topics); + CompletableFuture retCf = CompletableFuture.allOf(completeCfList.toArray(new CompletableFuture[0])); + completeCfList.clear(); + return retCf; + } + + private PartitionSnapshot snapshot(short funcVersion, Partition partition, + PartitionSnapshotVersion oldVersion, + PartitionSnapshotVersion newVersion, List> completeCfList) { + if (newVersion == null) { + // partition is closed + PartitionSnapshot snapshot = new PartitionSnapshot(); + snapshot.setPartitionIndex(partition.partitionId()); + snapshot.setLeaderEpoch(partition.getLeaderEpoch()); + snapshot.setOperation(SnapshotOperation.REMOVE.code()); + return snapshot; + } + return partition.withReadLock(() -> { + boolean includeSegments = oldVersion == null || oldVersion.segmentsVersion() < newVersion.segmentsVersion(); + PartitionSnapshot snapshot = new PartitionSnapshot(); + snapshot.setPartitionIndex(partition.partitionId()); + kafka.cluster.PartitionSnapshot src = partition.snapshot(); + completeCfList.add(src.completeCf()); + snapshot.setLeaderEpoch(src.leaderEpoch()); + SnapshotOperation operation = oldVersion == null ? SnapshotOperation.ADD : SnapshotOperation.PATCH; + snapshot.setOperation(operation.code()); + snapshot.setFirstUnstableOffset(logOffsetMetadata(src.firstUnstableOffset())); + snapshot.setLogEndOffset(logOffsetMetadata(src.logEndOffset())); + snapshot.setStreamMetadata(src.streamEndOffsets().entrySet() + .stream() + .map(e -> new AutomqGetPartitionSnapshotResponseData.StreamMetadata().setStreamId(e.getKey()).setEndOffset(e.getValue())) + .collect(Collectors.toList()) + ); + if (includeSegments) { + snapshot.setLogMetadata(logMetadata(src.logMeta())); + } + if (funcVersion > ZERO_ZONE_V0_REQUEST_VERSION) { + snapshot.setLastTimestampOffset(timestampOffset(src.lastTimestampOffset())); + } + return snapshot; + }); + } + + } + + static AutomqGetPartitionSnapshotResponseData.LogOffsetMetadata logOffsetMetadata(LogOffsetMetadata src) { + if (src == null) { + return null; + } + return new AutomqGetPartitionSnapshotResponseData.LogOffsetMetadata().setMessageOffset(src.messageOffset).setRelativePositionInSegment(src.relativePositionInSegment); + } + + static LogMetadata logMetadata(ElasticLogMeta src) { + if (src == null) { + return null; + } + LogMetadata logMetadata = new LogMetadata(); + + StreamMappingCollection streamMappingCollection = new StreamMappingCollection(); + src.getStreamMap().forEach((streamName, streamId) -> streamMappingCollection.add(new AutomqGetPartitionSnapshotResponseData.StreamMapping().setName(streamName).setStreamId(streamId))); + logMetadata.setStreamMap(streamMappingCollection); + + List segments = src.getSegmentMetas().stream().map(PartitionSnapshotsManager::segmentMetadata).collect(Collectors.toList()); + logMetadata.setSegments(segments); + + return logMetadata; + } + + static SegmentMetadata segmentMetadata(ElasticStreamSegmentMeta src) { + SegmentMetadata metadata = new SegmentMetadata(); + metadata.setBaseOffset(src.baseOffset()) + .setCreateTimestamp(src.createTimestamp()) + .setLastModifiedTimestamp(src.lastModifiedTimestamp()) + .setStreamSuffix(src.streamSuffix()) + .setLogSize(src.logSize()) + .setLog(sliceRange(src.log())) + .setTime(sliceRange(src.time())) + .setTransaction(sliceRange(src.txn())) + .setFirstBatchTimestamp(src.firstBatchTimestamp()) + .setTimeIndexLastEntry(timestampOffset(src.timeIndexLastEntry())); + return metadata; + } + + static AutomqGetPartitionSnapshotResponseData.SliceRange sliceRange(SliceRange src) { + return new AutomqGetPartitionSnapshotResponseData.SliceRange().setStart(src.start()).setEnd(src.end()); + } + + static AutomqGetPartitionSnapshotResponseData.TimestampOffsetData timestampOffset( + ElasticStreamSegmentMeta.TimestampOffsetData src) { + return new AutomqGetPartitionSnapshotResponseData.TimestampOffsetData().setTimestamp(src.timestamp()).setOffset(src.offset()); + } + + static AutomqGetPartitionSnapshotResponseData.TimestampOffsetData timestampOffset( + TimestampOffset src) { + return new AutomqGetPartitionSnapshotResponseData.TimestampOffsetData().setTimestamp(src.timestamp).setOffset(src.offset); + } + + static class PartitionWithVersion { + Partition partition; + PartitionSnapshotVersion version; + + public PartitionWithVersion(Partition partition, PartitionSnapshotVersion version) { + this.partition = partition; + this.version = version; + } + } + + static PartitionListener newPartitionListener(PartitionWithVersion version) { + return new PartitionListener() { + @Override + public void onNewLeaderEpoch(long oldEpoch, long newEpoch) { + version.version.incrementRecordsVersion(); + } + + @Override + public void onNewAppend(TopicPartition partition, long offset) { + version.version.incrementRecordsVersion(); + } + }; + } + + static LogEventListener newLogEventListener(PartitionWithVersion version) { + return (segment, event) -> version.version.incrementSegmentsVersion(); + } + +} diff --git a/core/src/main/java/kafka/automq/partition/snapshot/SnapshotOperation.java b/core/src/main/java/kafka/automq/partition/snapshot/SnapshotOperation.java new file mode 100644 index 0000000000..cd503e4202 --- /dev/null +++ b/core/src/main/java/kafka/automq/partition/snapshot/SnapshotOperation.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.partition.snapshot; + +public enum SnapshotOperation { + ADD((short) 0), PATCH((short) 1), REMOVE((short) 2); + + final short code; + + SnapshotOperation(short code) { + this.code = code; + } + + public short code() { + return code; + } + + public static SnapshotOperation parse(short code) { + switch (code) { + case 0: + return ADD; + case 1: + return PATCH; + case 2: + return REMOVE; + default: + throw new IllegalArgumentException("Unknown SnapshotOperation code: " + code); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/CatalogFactory.java b/core/src/main/java/kafka/automq/table/CatalogFactory.java new file mode 100644 index 0000000000..3706cc3370 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/CatalogFactory.java @@ -0,0 +1,257 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table; + +import kafka.server.KafkaConfig; + +import com.automq.stream.s3.operator.AwsObjectStorage; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.utils.IdURI; + +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.security.UserGroupInformation; +import org.apache.iceberg.CatalogUtil; +import org.apache.iceberg.catalog.Catalog; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.security.PrivilegedAction; +import java.util.Base64; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.function.Supplier; + +import static java.nio.file.StandardOpenOption.CREATE; +import static java.nio.file.StandardOpenOption.TRUNCATE_EXISTING; + +public class CatalogFactory { + private static final Logger LOGGER = LoggerFactory.getLogger(CatalogFactory.class); + private static final String CATALOG_TYPE_CONFIG = "type"; + + public Catalog newCatalog(KafkaConfig config) { + return new Builder(config).build(); + } + + static class Builder { + final KafkaConfig config; + + String catalogImpl; + BucketURI bucketURI; + final Map catalogConfigs; + final Map hadoopConfigs; + final Map options = new HashMap<>(); + Object hadoopConf = null; + UserGroupInformation ugi = null; + Catalog catalog = null; + + Builder(KafkaConfig config) { + this.config = config; + catalogConfigs = config.originalsWithPrefix("automq.table.topic.catalog."); + hadoopConfigs = config.originalsWithPrefix("automq.table.topic.hadoop."); + String catalogType = Optional.ofNullable(catalogConfigs.get(CATALOG_TYPE_CONFIG)).map(Object::toString).orElse(null); + if (StringUtils.isBlank(catalogType)) { + return; + } + bucketURI = config.automq().dataBuckets().get(0); + CredentialProviderHolder.setup(bucketURI); + options.put("ref", "main"); + options.put("client.region", bucketURI.region()); + options.put("client.credentials-provider", "kafka.automq.table.CredentialProviderHolder"); + switch (catalogType) { + case "glue": + withGlue(); + break; + case "nessie": + withNessie(); + break; + case "tablebucket": + withTableBucket(); + break; + case "hive": + withHive(); + break; + case "rest": + withRest(); + break; + default: + throw new IllegalArgumentException("Unsupported catalog type: " + catalogType); + } + catalogConfigs.forEach((k, v) -> options.put(k, v.toString())); + hadoopConf = mergeHadoopConfig(hadoopConfigs, hadoopConf); + options.remove(CATALOG_TYPE_CONFIG); + LOGGER.info("[TABLE_MANAGER_START],catalog={},options={},hadoopConfig={}", catalogType, options, hadoopConf); + this.catalog = runAs(() -> CatalogUtil.loadCatalog(catalogImpl, catalogType, options, hadoopConf)); + } + + private Object mergeHadoopConfig(Map from, Object hadoopConf) { + if (!from.isEmpty()) { + if (hadoopConf == null) { + hadoopConf = new Configuration(); + } + for (Map.Entry entry : hadoopConfigs.entrySet()) { + String k = entry.getKey(); + Object v = entry.getValue(); + ((Configuration) hadoopConf).set(k, v.toString()); + } + } + return hadoopConf; + } + + public Catalog build() { + return catalog; + } + + private void withGlue() { + catalogImpl = "org.apache.iceberg.aws.glue.GlueCatalog"; + if (StringUtils.isNotBlank(bucketURI.endpoint())) { + options.put("glue.endpoint", bucketURI.endpoint().replaceFirst("s3", "glue")); + } + putDataBucketAsWarehouse(false); + } + + private void withNessie() { + // nessie config extension e.g. + // automq.table.topic.catalog.uri=http://localhost:19120/api/v2 + catalogImpl = "org.apache.iceberg.nessie.NessieCatalog"; + putDataBucketAsWarehouse(false); + } + + private void withTableBucket() { + // table bucket config extension e.g. + // automq.table.topic.catalog.warehouse=table bucket arn + catalogImpl = "software.amazon.s3tables.iceberg.S3TablesCatalog"; + } + + private void withHive() { + // hive config extension e.g. + // automq.table.topic.catalog.uri=thrift://xxx:9083 + // kerberos authentication + // - automq.table.topic.catalog.auth=kerberos://?principal=base64(clientPrincipal)&keytab=base64(keytabFile)&krb5conf=base64(krb5confFile) + // - automq.table.topic.hadoop.metastore.kerberos.principal=serverPrincipal + + // simple authentication + // - automq.table.topic.catalog.auth=simple://?username=xxx + catalogImpl = "org.apache.iceberg.hive.HiveCatalog"; + putDataBucketAsWarehouse(true); + + hadoopConf = new Configuration(); + IdURI uri = IdURI.parse("0@" + catalogConfigs.getOrDefault("auth", "none://?")); + try { + switch (uri.protocol()) { + case "kerberos": { + System.setProperty("sun.security.krb5.debug", "true"); + String configBasePath = config.metadataLogDir(); + System.setProperty( + "java.security.krb5.conf", + base64Config2file(uri.extensionString("krb5conf"), configBasePath, "krb5.conf") + ); + Configuration configuration = new Configuration(); + configuration.set("hadoop.security.authentication", "Kerberos"); + UserGroupInformation.setConfiguration(configuration); + UserGroupInformation.loginUserFromKeytab( + decodeBase64(uri.extensionString("principal")), + base64Config2file(uri.extensionString("keytab"), configBasePath, "keytab") + ); + ugi = UserGroupInformation.getCurrentUser(); + ((Configuration) hadoopConf).set("metastore.sasl.enabled", "true"); + break; + } + case "simple": { + ugi = UserGroupInformation.createRemoteUser(uri.extensionString("username")); + UserGroupInformation.setLoginUser(ugi); + ((Configuration) hadoopConf).set("metastore.sasl.enabled", "true"); + break; + } + default: { + } + } + } catch (Throwable e) { + throw new RuntimeException(e); + } + } + + private void withRest() { + // rest config extension e.g. + // automq.table.topic.catalog.uri=http://127.0.0.1:9001/iceberg + // If a token is set, HTTP requests use the value as a bearer token in the HTTP Authorization header. + // If credential is used, then the key and secret are used to fetch a token using the OAuth2 client credentials flow. + // The resulting token is used as the bearer token for subsequent requests. + // config ref. org.apache.iceberg.rest.RESTSessionCatalog#initialize + // automq.table.topic.catalog.oauth2-server-uri= + // automq.table.topic.catalog.credential= + // automq.table.topic.catalog.token= + // automq.table.topic.catalog.scope= + catalogImpl = "org.apache.iceberg.rest.RESTCatalog"; + putDataBucketAsWarehouse(false); + } + + private Catalog runAs(Supplier func) { + if (ugi != null) { + return ugi.doAs((PrivilegedAction) func::get); + } else { + return func.get(); + } + } + + // important: use putIfAbsent to let the user override all values directly in catalog configuration + private void putDataBucketAsWarehouse(boolean s3a) { + if (StringUtils.isNotBlank(bucketURI.endpoint())) { + options.putIfAbsent("s3.endpoint", bucketURI.endpoint()); + } + if (bucketURI.extensionBool(AwsObjectStorage.PATH_STYLE_KEY, false)) { + options.putIfAbsent("s3.path-style-access", "true"); + } + options.putIfAbsent("io-impl", "org.apache.iceberg.aws.s3.S3FileIO"); + options.putIfAbsent("warehouse", String.format((s3a ? "s3a" : "s3") + "://%s/iceberg", bucketURI.bucket())); + } + + } + + /** + * Decode base64 str and save it to file + * + * @return the file path + */ + private static String base64Config2file(String base64, String configPath, String configName) { + byte[] bytes = Base64.getDecoder().decode(base64); + try { + Path dir = Paths.get(configPath); + if (!Files.exists(dir)) { + Files.createDirectories(dir); + } + Path filePath = Paths.get(configPath + File.separator + configName); + Files.write(filePath, bytes, CREATE, TRUNCATE_EXISTING); + return filePath.toAbsolutePath().toString(); + } catch (Throwable e) { + throw new RuntimeException(e); + } + } + + private static String decodeBase64(String base64) { + return new String(Base64.getDecoder().decode(base64), StandardCharsets.ISO_8859_1); + } +} diff --git a/core/src/main/java/kafka/automq/table/Channel.java b/core/src/main/java/kafka/automq/table/Channel.java new file mode 100644 index 0000000000..a0b1225811 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/Channel.java @@ -0,0 +1,186 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table; + +import kafka.automq.table.events.AvroCodec; +import kafka.automq.table.events.Envelope; +import kafka.automq.table.events.Event; +import kafka.automq.utils.ClientUtils; +import kafka.server.KafkaConfig; + +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.ConsumerRecords; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.clients.producer.Callback; +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.clients.producer.RecordMetadata; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.ByteArraySerializer; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.time.Duration; +import java.util.List; +import java.util.Properties; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.LinkedBlockingQueue; + +public class Channel { + public static final long LATEST_OFFSET = -1L; + private static final Logger LOGGER = LoggerFactory.getLogger(Channel.class); + private static final String CONTROL_TOPIC = "__automq_table_control"; + private static final String DATA_TOPIC = "__automq_table_data"; + private final Properties clientBaseConfigs; + private final KafkaProducer producer; + private volatile int dataTopicPartitionNums = -1; + + public Channel(KafkaConfig kafkaConfig) { + clientBaseConfigs = ClientUtils.clusterClientBaseConfig(kafkaConfig); + Properties props = new Properties(); + props.putAll(clientBaseConfigs); + props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + producer = new KafkaProducer<>(props); + } + + public void send(String topic, Event event) throws Exception { + asyncSend(topic, event).get(); + } + + public CompletableFuture asyncSend(String topic, Event event) { + CompletableFuture cf = new CompletableFuture<>(); + try { + Callback callback = new Callback() { + @Override + public void onCompletion(RecordMetadata metadata, Exception e) { + if (e != null) { + cf.completeExceptionally(e); + } else { + cf.complete(null); + } + } + }; + switch (event.type()) { + case COMMIT_REQUEST: + producer.send(new ProducerRecord<>(CONTROL_TOPIC, 0, null, AvroCodec.encode(event)), callback); + break; + case COMMIT_RESPONSE: + int partitionId = Math.abs(topic.hashCode() % dataTopicPartitionNums()); + producer.send(new ProducerRecord<>(DATA_TOPIC, partitionId, null, AvroCodec.encode(event)), callback); + break; + default: + throw new IllegalArgumentException("Unknown event type: " + event.type()); + } + } catch (Throwable e) { + cf.completeExceptionally(e); + } + return cf; + } + + public SubChannel subscribeControl() { + Properties props = new Properties(); + props.putAll(clientBaseConfigs); + props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false); + props.put(ConsumerConfig.GROUP_ID_CONFIG, "automq-table-control-consumer"); + props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + @SuppressWarnings("resource") + KafkaConsumer consumer = new KafkaConsumer<>(props); + consumer.assign(List.of(new TopicPartition(CONTROL_TOPIC, 0))); + return new KafkaConsumerSubChannel(consumer); + } + + public SubChannel subscribeData(String topic, long offset) { + Properties props = new Properties(); + props.putAll(clientBaseConfigs); + props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false); + props.put(ConsumerConfig.GROUP_ID_CONFIG, "automq-table-data-consumer"); + props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + @SuppressWarnings("resource") + KafkaConsumer consumer = new KafkaConsumer<>(props); + int partitionNums = dataTopicPartitionNums(); + TopicPartition topicPartition = new TopicPartition(DATA_TOPIC, Math.abs(topic.hashCode() % partitionNums)); + consumer.assign(List.of(topicPartition)); + if (offset == LATEST_OFFSET) { + consumer.seekToEnd(List.of(topicPartition)); + } else { + consumer.seek(topicPartition, offset); + } + return new KafkaConsumerSubChannel(consumer); + } + + private int dataTopicPartitionNums() { + if (dataTopicPartitionNums > 0) { + return dataTopicPartitionNums; + } + synchronized (this) { + if (dataTopicPartitionNums > 0) { + return dataTopicPartitionNums; + } + dataTopicPartitionNums = producer.partitionsFor(DATA_TOPIC).size(); + return dataTopicPartitionNums; + } + } + + public interface SubChannel { + Envelope poll(); + + void close(); + } + + public static class KafkaConsumerSubChannel implements SubChannel { + private final KafkaConsumer consumer; + private final Queue left = new LinkedBlockingQueue<>(); + + public KafkaConsumerSubChannel(KafkaConsumer consumer) { + this.consumer = consumer; + } + + @Override + public Envelope poll() { + if (!left.isEmpty()) { + return left.poll(); + } + ConsumerRecords records = consumer.poll(Duration.ofSeconds(0)); + records.forEach(record -> { + try { + Event event = AvroCodec.decode(record.value()); + left.add(new Envelope(record.partition(), record.offset(), event)); + } catch (IOException e) { + LOGGER.error("decode fail"); + } + }); + return left.poll(); + } + + @Override + public void close() { + consumer.close(); + } + } + +} diff --git a/core/src/main/java/kafka/automq/table/CredentialProviderHolder.java b/core/src/main/java/kafka/automq/table/CredentialProviderHolder.java new file mode 100644 index 0000000000..f0264e7632 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/CredentialProviderHolder.java @@ -0,0 +1,75 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table; + +import com.automq.stream.s3.operator.AutoMQStaticCredentialsProvider; +import com.automq.stream.s3.operator.BucketURI; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; + +import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProviderChain; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; + + +public class CredentialProviderHolder implements AwsCredentialsProvider { + private static Function providerSupplier = bucketURI -> newCredentialsProviderChain( + credentialsProviders(bucketURI)); + private static BucketURI bucketURI; + + public static void setup(Function providerSupplier) { + CredentialProviderHolder.providerSupplier = providerSupplier; + } + + public static void setup(BucketURI bucketURI) { + CredentialProviderHolder.bucketURI = bucketURI; + } + + private static List credentialsProviders(BucketURI bucketURI) { + return List.of(new AutoMQStaticCredentialsProvider(bucketURI), DefaultCredentialsProvider.builder().build()); + } + + private static AwsCredentialsProvider newCredentialsProviderChain( + List credentialsProviders) { + List providers = new ArrayList<>(credentialsProviders); + providers.add(AnonymousCredentialsProvider.create()); + return AwsCredentialsProviderChain.builder() + .reuseLastProviderEnabled(true) + .credentialsProviders(providers) + .build(); + } + + // iceberg will invoke create with reflection. + public static AwsCredentialsProvider create() { + if (bucketURI == null) { + throw new IllegalStateException("BucketURI must be set before calling create(). Please invoke setup(BucketURI) first."); + } + return providerSupplier.apply(bucketURI); + } + + @Override + public AwsCredentials resolveCredentials() { + throw new UnsupportedOperationException(); + } +} diff --git a/core/src/main/java/kafka/automq/table/TableManager.java b/core/src/main/java/kafka/automq/table/TableManager.java new file mode 100644 index 0000000000..b75e8ab05b --- /dev/null +++ b/core/src/main/java/kafka/automq/table/TableManager.java @@ -0,0 +1,151 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table; + +import kafka.automq.table.coordinator.TableCoordinator; +import kafka.automq.table.worker.TableWorkers; +import kafka.cluster.Partition; +import kafka.log.streamaspect.ElasticLog; +import kafka.log.streamaspect.ElasticUnifiedLog; +import kafka.log.streamaspect.MetaStream; +import kafka.server.KafkaConfig; +import kafka.server.MetadataCache; +import kafka.server.streamaspect.PartitionLifecycleListener; + +import org.apache.kafka.common.TopicPartition; + +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.threads.EventLoop; + +import org.apache.iceberg.catalog.Catalog; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +public class TableManager implements PartitionLifecycleListener { + private static final Logger LOGGER = LoggerFactory.getLogger(TableManager.class); + private final Catalog catalog; + private final Channel channel; + private final EventLoop[] coordinatorEventLoops; + private final Map coordinators = new HashMap<>(); + private final Set tableTopicPartitions = new HashSet<>(); + private final TableWorkers tableWorkers; + private final KafkaConfig config; + + private final MetadataCache metadataCache; + + public TableManager(MetadataCache metadataCache, KafkaConfig config) { + this.metadataCache = metadataCache; + this.config = config; + this.catalog = new CatalogFactory().newCatalog(config); + if (this.catalog == null) { + channel = null; + coordinatorEventLoops = null; + tableWorkers = null; + return; + } + this.channel = new Channel(config); + this.tableWorkers = new TableWorkers(catalog, channel, config); + this.coordinatorEventLoops = new EventLoop[Math.max(Systems.CPU_CORES / 2, 1)]; + for (int i = 0; i < coordinatorEventLoops.length; i++) { + this.coordinatorEventLoops[i] = new EventLoop("table-coordinator-" + i); + } + } + + @Override + public synchronized void onOpen(Partition partition) { + if (catalog == null) { + return; + } + ElasticUnifiedLog log = (ElasticUnifiedLog) partition.log().get(); + log.addConfigChangeListener((l, config) -> { + synchronized (TableManager.this) { + if (config.tableTopicEnable && !tableTopicPartitions.contains(partition.topicPartition())) { + add(partition); + } else if (!config.tableTopicEnable && tableTopicPartitions.contains(partition.topicPartition())) { + remove(partition); + } + } + }); + if (log.config().tableTopicEnable && !tableTopicPartitions.contains(partition.topicPartition())) { + add(partition); + } + } + + private synchronized void add(Partition partition) { + try { + String topic = partition.topicPartition().topic(); + int partitionId = partition.topicPartition().partition(); + if (partitionId == 0) { + // start coordinator + EventLoop eventLoop = coordinatorEventLoops[Math.abs(topic.hashCode() % coordinatorEventLoops.length)]; + MetaStream metaStream = ((ElasticLog) (partition.log().get().localLog())).metaStream(); + //noinspection resource + TableCoordinator coordinator = new TableCoordinator(catalog, topic, metaStream, channel, eventLoop, metadataCache, () -> partition.log().get().config()); + coordinators.put(partition.topicPartition(), coordinator); + coordinator.start(); + } + // start worker + tableWorkers.add(partition); + tableTopicPartitions.add(partition.topicPartition()); + } catch (Throwable e) { + LOGGER.error("[TABLE_TOPIC_PARTITION_ADD],{}", partition.topicPartition(), e); + } + } + + @Override + public void onClose(Partition partition) { + remove(partition); + } + + private synchronized void remove(Partition partition) { + if (catalog == null) { + return; + } + try { + int partitionId = partition.topicPartition().partition(); + if (partitionId == 0) { + TableCoordinator coordinator = coordinators.remove(partition.topicPartition()); + if (coordinator != null) { + coordinator.close(); + } + } + tableWorkers.remove(partition); + tableTopicPartitions.remove(partition.topicPartition()); + } catch (Throwable e) { + LOGGER.error("[TABLE_TOPIC_PARTITION_DELETE],{}", partition.topicPartition(), e); + } + } + + public void close() { + if (catalog == null) { + return; + } + try { + tableWorkers.close(); + } catch (Throwable e) { + LOGGER.error("[TABLE_MANAGER_CLOSE_FAIL]", e); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/binder/AbstractTypeAdapter.java b/core/src/main/java/kafka/automq/table/binder/AbstractTypeAdapter.java new file mode 100644 index 0000000000..18cfac5899 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/binder/AbstractTypeAdapter.java @@ -0,0 +1,229 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.binder; + +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.ByteBuffers; +import org.apache.iceberg.util.DateTimeUtil; +import org.apache.iceberg.util.UUIDUtil; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalTime; +import java.time.ZoneOffset; +import java.time.temporal.Temporal; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * Abstract implementation providing common type conversion logic from source formats + * to Iceberg's internal Java type representation. + *

+ * Handles dispatch logic and provides default conversion implementations for primitive types. + * Subclasses implement format-specific conversion for complex types (LIST, MAP, STRUCT). + * + * @param The type of the source schema (e.g., org.apache.avro.Schema) + */ +public abstract class AbstractTypeAdapter implements TypeAdapter { + + + @SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"}) + @Override + public Object convert(Object sourceValue, S sourceSchema, Type targetType, StructConverter structConverter) { + if (sourceValue == null) { + return null; + } + + switch (targetType.typeId()) { + case BOOLEAN: + return convertBoolean(sourceValue, sourceSchema, targetType); + case INTEGER: + return convertInteger(sourceValue, sourceSchema, targetType); + case LONG: + return convertLong(sourceValue, sourceSchema, targetType); + case FLOAT: + return convertFloat(sourceValue, sourceSchema, targetType); + case DOUBLE: + return convertDouble(sourceValue, sourceSchema, targetType); + case STRING: + return convertString(sourceValue, sourceSchema, targetType); + case BINARY: + return convertBinary(sourceValue, sourceSchema, targetType); + case FIXED: + return convertFixed(sourceValue, sourceSchema, targetType); + case UUID: + return convertUUID(sourceValue, sourceSchema, targetType); + case DECIMAL: + return convertDecimal(sourceValue, sourceSchema, (Types.DecimalType) targetType); + case DATE: + return convertDate(sourceValue, sourceSchema, targetType); + case TIME: + return convertTime(sourceValue, sourceSchema, targetType); + case TIMESTAMP: + return convertTimestamp(sourceValue, sourceSchema, (Types.TimestampType) targetType); + case LIST: + return convertList(sourceValue, sourceSchema, (Types.ListType) targetType, structConverter); + case MAP: + return convertMap(sourceValue, sourceSchema, (Types.MapType) targetType, structConverter); + case STRUCT: + return structConverter.convert(sourceValue, sourceSchema, targetType); + default: + return sourceValue; + } + } + + protected Object convertBoolean(Object sourceValue, S ignoredSourceSchema, Type targetType) { + if (sourceValue instanceof Boolean) return sourceValue; + if (sourceValue instanceof String) return Boolean.parseBoolean((String) sourceValue); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertInteger(Object sourceValue, S ignoredSourceSchema, Type targetType) { + if (sourceValue instanceof Integer) return sourceValue; + if (sourceValue instanceof Number) return ((Number) sourceValue).intValue(); + if (sourceValue instanceof String) return Integer.parseInt((String) sourceValue); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertLong(Object sourceValue, S ignoredSourceSchema, Type targetType) { + if (sourceValue instanceof Long) return sourceValue; + if (sourceValue instanceof Number) return ((Number) sourceValue).longValue(); + if (sourceValue instanceof String) return Long.parseLong((String) sourceValue); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertFloat(Object sourceValue, S ignoredSourceSchema, Type targetType) { + if (sourceValue instanceof Float) return sourceValue; + if (sourceValue instanceof Number) return ((Number) sourceValue).floatValue(); + if (sourceValue instanceof String) return Float.parseFloat((String) sourceValue); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertDouble(Object sourceValue, S ignoredSourceSchema, Type targetType) { + if (sourceValue instanceof Double) return sourceValue; + if (sourceValue instanceof Number) return ((Number) sourceValue).doubleValue(); + if (sourceValue instanceof String) return Double.parseDouble((String) sourceValue); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertString(Object sourceValue, S sourceSchema, Type targetType) { + if (sourceValue instanceof String) { + return sourceValue; + } + // Simple toString conversion - subclasses can override for more complex logic + return sourceValue.toString(); + } + + protected Object convertBinary(Object sourceValue, S sourceSchema, Type targetType) { + if (sourceValue instanceof ByteBuffer) return ((ByteBuffer) sourceValue).duplicate(); + if (sourceValue instanceof byte[]) return ByteBuffer.wrap((byte[]) sourceValue); + if (sourceValue instanceof String) return ByteBuffer.wrap(((String) sourceValue).getBytes(StandardCharsets.UTF_8)); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertFixed(Object sourceValue, S sourceSchema, Type targetType) { + if (sourceValue instanceof byte[]) return sourceValue; + if (sourceValue instanceof ByteBuffer) return ByteBuffers.toByteArray((ByteBuffer) sourceValue); + if (sourceValue instanceof String) return ((String) sourceValue).getBytes(StandardCharsets.UTF_8); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertUUID(Object sourceValue, S sourceSchema, Type targetType) { + UUID uuid = null; + if (sourceValue instanceof String) { + uuid = UUID.fromString(sourceValue.toString()); + } else if (sourceValue instanceof UUID) { + uuid = (UUID) sourceValue; + } else if (sourceValue instanceof ByteBuffer) { + ByteBuffer bb = ((ByteBuffer) sourceValue).duplicate(); + if (bb.remaining() == 16) { + uuid = new UUID(bb.getLong(), bb.getLong()); + } + } + if (uuid != null) { + return UUIDUtil.convert(uuid); + } + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertDecimal(Object sourceValue, S ignoredSourceSchema, Types.DecimalType targetType) { + if (sourceValue instanceof BigDecimal) return sourceValue; + if (sourceValue instanceof String) return new BigDecimal((String) sourceValue); + if (sourceValue instanceof byte[]) return new BigDecimal(new java.math.BigInteger((byte[]) sourceValue), targetType.scale()); + if (sourceValue instanceof ByteBuffer) { + ByteBuffer bb = ((ByteBuffer) sourceValue).duplicate(); + byte[] bytes = new byte[bb.remaining()]; + bb.get(bytes); + return new BigDecimal(new java.math.BigInteger(bytes), targetType.scale()); + } + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertDate(Object sourceValue, S ignoredSourceSchema, Type targetType) { + if (sourceValue instanceof LocalDate) return sourceValue; + if (sourceValue instanceof Number) return LocalDate.ofEpochDay(((Number) sourceValue).intValue()); + if (sourceValue instanceof Date) return ((Date) sourceValue).toInstant().atZone(ZoneOffset.UTC).toLocalDate(); + if (sourceValue instanceof String) return LocalDate.parse(sourceValue.toString()); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertTime(Object sourceValue, S sourceSchema, Type targetType) { + if (sourceValue instanceof LocalTime) return sourceValue; + if (sourceValue instanceof Date) return ((Date) sourceValue).toInstant().atZone(ZoneOffset.UTC).toLocalTime(); + if (sourceValue instanceof String) return LocalTime.parse(sourceValue.toString()); + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected Object convertTimestamp(Object sourceValue, S sourceSchema, Types.TimestampType targetType) { + if (sourceValue instanceof Temporal) return sourceValue; + if (sourceValue instanceof Date) { + Instant instant = ((Date) sourceValue).toInstant(); + long micros = DateTimeUtil.microsFromInstant(instant); + return targetType.shouldAdjustToUTC() + ? DateTimeUtil.timestamptzFromMicros(micros) + : DateTimeUtil.timestampFromMicros(micros); + } + if (sourceValue instanceof String) { + Instant instant = Instant.parse(sourceValue.toString()); + long micros = DateTimeUtil.microsFromInstant(instant); + return targetType.shouldAdjustToUTC() + ? DateTimeUtil.timestamptzFromMicros(micros) + : DateTimeUtil.timestampFromMicros(micros); + } + if (sourceValue instanceof Number) { + // Assume the number represents microseconds since epoch + // Subclasses should override to handle milliseconds or other units based on logical type + long micros = ((Number) sourceValue).longValue(); + return targetType.shouldAdjustToUTC() + ? DateTimeUtil.timestamptzFromMicros(micros) + : DateTimeUtil.timestampFromMicros(micros); + } + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to " + targetType.typeId()); + } + + protected abstract List convertList(Object sourceValue, S sourceSchema, Types.ListType targetType, StructConverter structConverter); + + protected abstract Map convertMap(Object sourceValue, S sourceSchema, Types.MapType targetType, StructConverter structConverter); +} diff --git a/core/src/main/java/kafka/automq/table/binder/AvroValueAdapter.java b/core/src/main/java/kafka/automq/table/binder/AvroValueAdapter.java new file mode 100644 index 0000000000..d7eae3a19b --- /dev/null +++ b/core/src/main/java/kafka/automq/table/binder/AvroValueAdapter.java @@ -0,0 +1,210 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.binder; + +import org.apache.avro.LogicalType; +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.util.Utf8; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.DateTimeUtil; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * A concrete implementation of TypeAdapter that converts values from Avro's + * data representation to Iceberg's internal Java type representation. + *

+ * This class extends {@link AbstractTypeAdapter} and overrides methods to handle + * Avro-specific types like Utf8, EnumSymbol, and Fixed, as well as Avro's + * specific representations for List and Map. + */ +public class AvroValueAdapter extends AbstractTypeAdapter { + private static final org.apache.avro.Schema STRING_SCHEMA_INSTANCE = org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING); + + @Override + protected Object convertString(Object sourceValue, Schema sourceSchema, Type targetType) { + if (sourceValue instanceof Utf8) { + return sourceValue; + } + if (sourceValue instanceof GenericData.EnumSymbol) { + return sourceValue.toString(); + } + return super.convertString(sourceValue, sourceSchema, targetType); + } + + @Override + protected Object convertBinary(Object sourceValue, Schema sourceSchema, Type targetType) { + if (sourceValue instanceof GenericData.Fixed) { + return ByteBuffer.wrap(((GenericData.Fixed) sourceValue).bytes()); + } + return super.convertBinary(sourceValue, sourceSchema, targetType); + } + + @Override + protected Object convertFixed(Object sourceValue, Schema sourceSchema, Type targetType) { + if (sourceValue instanceof GenericData.Fixed) { + return ((GenericData.Fixed) sourceValue).bytes(); + } + return super.convertFixed(sourceValue, sourceSchema, targetType); + } + + @Override + protected Object convertUUID(Object sourceValue, Schema sourceSchema, Type targetType) { + if (sourceValue instanceof Utf8) { + return super.convertUUID(sourceValue.toString(), sourceSchema, targetType); + } + return super.convertUUID(sourceValue, sourceSchema, targetType); + } + + @Override + protected Object convertTime(Object sourceValue, Schema sourceSchema, Type targetType) { + if (sourceValue instanceof Number) { + LogicalType logicalType = sourceSchema.getLogicalType(); + if (logicalType instanceof LogicalTypes.TimeMicros) { + return DateTimeUtil.timeFromMicros(((Number) sourceValue).longValue()); + } else if (logicalType instanceof LogicalTypes.TimeMillis) { + return DateTimeUtil.timeFromMicros(((Number) sourceValue).longValue() * 1000); + } + } + return super.convertTime(sourceValue, sourceSchema, targetType); + } + + @Override + protected Object convertTimestamp(Object sourceValue, Schema sourceSchema, Types.TimestampType targetType) { + if (sourceValue instanceof Number) { + long value = ((Number) sourceValue).longValue(); + LogicalType logicalType = sourceSchema.getLogicalType(); + if (logicalType instanceof LogicalTypes.TimestampMillis) { + return targetType.shouldAdjustToUTC() + ? DateTimeUtil.timestamptzFromMicros(value * 1000) + : DateTimeUtil.timestampFromMicros(value * 1000); + } else if (logicalType instanceof LogicalTypes.TimestampMicros) { + return targetType.shouldAdjustToUTC() + ? DateTimeUtil.timestamptzFromMicros(value) + : DateTimeUtil.timestampFromMicros(value); + } else if (logicalType instanceof LogicalTypes.LocalTimestampMillis) { + return DateTimeUtil.timestampFromMicros(value * 1000); + } else if (logicalType instanceof LogicalTypes.LocalTimestampMicros) { + return DateTimeUtil.timestampFromMicros(value); + } + } + return super.convertTimestamp(sourceValue, sourceSchema, targetType); + } + + @Override + protected List convertList(Object sourceValue, Schema sourceSchema, Types.ListType targetType, StructConverter structConverter) { + Schema listSchema = sourceSchema; + Schema elementSchema = listSchema.getElementType(); + + List sourceList; + if (sourceValue instanceof GenericData.Array) { + sourceList = (GenericData.Array) sourceValue; + } else if (sourceValue instanceof List) { + sourceList = (List) sourceValue; + } else { + throw new IllegalArgumentException("Cannot convert " + sourceValue.getClass().getSimpleName() + " to LIST"); + } + + List list = new ArrayList<>(sourceList.size()); + for (Object element : sourceList) { + Object convert = convert(element, elementSchema, targetType.elementType(), structConverter); + list.add(convert); + } + return list; + } + + @Override + protected Map convertMap(Object sourceValue, Schema sourceSchema, Types.MapType targetType, StructConverter structConverter) { + if (sourceValue instanceof GenericData.Array) { + GenericData.Array arrayValue = (GenericData.Array) sourceValue; + Map recordMap = new HashMap<>(arrayValue.size()); + + Schema kvSchema = sourceSchema.getElementType(); + + Schema.Field keyField = kvSchema.getFields().get(0); + Schema.Field valueField = kvSchema.getFields().get(1); + if (keyField == null || valueField == null) { + throw new IllegalStateException("Map entry schema missing key/value fields: " + kvSchema); + } + + Schema keySchema = keyField.schema(); + Schema valueSchema = valueField.schema(); + Type keyType = targetType.keyType(); + Type valueType = targetType.valueType(); + + for (Object element : arrayValue) { + if (element == null) { + continue; + } + GenericRecord record = (GenericRecord) element; + Object key = convert(record.get(keyField.pos()), keySchema, keyType, structConverter); + Object value = convert(record.get(valueField.pos()), valueSchema, valueType, structConverter); + recordMap.put(key, value); + } + return recordMap; + } + + Schema mapSchema = sourceSchema; + + Map sourceMap = (Map) sourceValue; + Map adaptedMap = new HashMap<>(sourceMap.size()); + + Schema valueSchema = mapSchema.getValueType(); + Type keyType = targetType.keyType(); + Type valueType = targetType.valueType(); + + for (Map.Entry entry : sourceMap.entrySet()) { + Object rawKey = entry.getKey(); + Object key = convert(rawKey, STRING_SCHEMA_INSTANCE, keyType, structConverter); + Object value = convert(entry.getValue(), valueSchema, valueType, structConverter); + adaptedMap.put(key, value); + } + return adaptedMap; + } + + @Override + public Object convert(Object sourceValue, Schema sourceSchema, Type targetType) { + return convert(sourceValue, sourceSchema, targetType, this::convertStruct); + } + + protected Object convertStruct(Object sourceValue, Schema sourceSchema, Type targetType) { + org.apache.iceberg.Schema schema = targetType.asStructType().asSchema(); + org.apache.iceberg.data.GenericRecord result = org.apache.iceberg.data.GenericRecord.create(schema); + for (Types.NestedField f : schema.columns()) { + // Convert the value to the expected type + GenericRecord record = (GenericRecord) sourceValue; + Schema.Field sourceField = sourceSchema.getField(f.name()); + if (sourceField == null) { + throw new IllegalStateException("Missing field '" + f.name() + + "' in source schema: " + sourceSchema.getFullName()); + } + Object fieldValue = convert(record.get(f.name()), sourceField.schema(), f.type()); + result.setField(f.name(), fieldValue); + } + return result; + } +} diff --git a/core/src/main/java/kafka/automq/table/binder/FieldMapping.java b/core/src/main/java/kafka/automq/table/binder/FieldMapping.java new file mode 100644 index 0000000000..13178235f4 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/binder/FieldMapping.java @@ -0,0 +1,57 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.binder; + +import org.apache.avro.Schema; +import org.apache.iceberg.types.Type; + +/** + * Represents the mapping between an Avro field and its corresponding Iceberg field. + * This class stores the position, key, schema, and type information needed to + * convert field values during record binding. + */ +public class FieldMapping { + private final int avroPosition; + private final String avroKey; + private final Type icebergType; + private final Schema avroSchema; + + public FieldMapping(int avroPosition, String avroKey, Type icebergType, Schema avroSchema) { + this.avroPosition = avroPosition; + this.avroKey = avroKey; + this.icebergType = icebergType; + this.avroSchema = avroSchema; + } + + public int avroPosition() { + return avroPosition; + } + + public String avroKey() { + return avroKey; + } + + public Type icebergType() { + return icebergType; + } + + public Schema avroSchema() { + return avroSchema; + } +} diff --git a/core/src/main/java/kafka/automq/table/binder/RecordBinder.java b/core/src/main/java/kafka/automq/table/binder/RecordBinder.java new file mode 100644 index 0000000000..47835c0355 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/binder/RecordBinder.java @@ -0,0 +1,494 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.binder; + + +import kafka.automq.table.metric.FieldMetric; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +import static org.apache.avro.Schema.Type.ARRAY; +import static org.apache.avro.Schema.Type.NULL; + +/** + * A factory that creates lazy-evaluation Record views of Avro GenericRecords. + * Field values are converted only when accessed, avoiding upfront conversion overhead. + */ +public class RecordBinder { + + private final org.apache.iceberg.Schema icebergSchema; + private final TypeAdapter typeAdapter; + private final Map fieldNameToPosition; + private final FieldMapping[] fieldMappings; + + // Pre-computed RecordBinders for nested STRUCT fields + private final Map nestedStructBinders; + + // Field count statistics for this batch + private final AtomicLong batchFieldCount; + + + public RecordBinder(GenericRecord avroRecord) { + this(AvroSchemaUtil.toIceberg(avroRecord.getSchema()), avroRecord.getSchema()); + } + + public RecordBinder(org.apache.iceberg.Schema icebergSchema, Schema avroSchema) { + this(icebergSchema, avroSchema, new AvroValueAdapter()); + } + + public RecordBinder(org.apache.iceberg.Schema icebergSchema, Schema avroSchema, TypeAdapter typeAdapter) { + this(icebergSchema, avroSchema, typeAdapter, new AtomicLong(0)); + } + + public RecordBinder(org.apache.iceberg.Schema icebergSchema, Schema avroSchema, TypeAdapter typeAdapter, AtomicLong batchFieldCount) { + this.icebergSchema = icebergSchema; + this.typeAdapter = typeAdapter; + this.batchFieldCount = batchFieldCount; + + // Pre-compute field name to position mapping + this.fieldNameToPosition = new HashMap<>(); + for (int i = 0; i < icebergSchema.columns().size(); i++) { + fieldNameToPosition.put(icebergSchema.columns().get(i).name(), i); + } + + // Initialize field mappings + this.fieldMappings = buildFieldMappings(avroSchema, icebergSchema); + // Pre-compute nested struct binders + this.nestedStructBinders = precomputeBindersMap(typeAdapter); + } + + public RecordBinder createBinderForNewSchema(org.apache.iceberg.Schema icebergSchema, Schema avroSchema) { + return new RecordBinder(icebergSchema, avroSchema, typeAdapter, batchFieldCount); + } + + + public org.apache.iceberg.Schema getIcebergSchema() { + return icebergSchema; + } + + /** + * Creates a new immutable Record view of the given Avro record. + * Each call returns a separate instance with its own data reference. + */ + public Record bind(GenericRecord avroRecord) { + if (avroRecord == null) { + return null; + } + return new AvroRecordView(avroRecord, icebergSchema, typeAdapter, + fieldNameToPosition, fieldMappings, nestedStructBinders, this); + } + + /** + * Gets the accumulated field count for this batch and resets it to zero. + * Should be called after each flush to collect field statistics. + */ + public long getAndResetFieldCount() { + return batchFieldCount.getAndSet(0); + } + + /** + * Adds field count to the batch total. Called by AvroRecordView instances. + */ + void addFieldCount(long count) { + batchFieldCount.addAndGet(count); + } + + private FieldMapping[] buildFieldMappings(Schema avroSchema, org.apache.iceberg.Schema icebergSchema) { + Schema recordSchema = avroSchema; + FieldMapping[] mappings = new FieldMapping[icebergSchema.columns().size()]; + + // Unwrap UNION if it contains only one non-NULL type + recordSchema = resolveUnionElement(recordSchema); + + for (int icebergPos = 0; icebergPos < icebergSchema.columns().size(); icebergPos++) { + Types.NestedField icebergField = icebergSchema.columns().get(icebergPos); + String fieldName = icebergField.name(); + + Schema.Field avroField = recordSchema.getField(fieldName); + if (avroField != null) { + mappings[icebergPos] = buildFieldMapping( + avroField.name(), + avroField.pos(), + icebergField.type(), + avroField.schema() + ); + } else { + mappings[icebergPos] = null; + } + } + return mappings; + } + + private FieldMapping buildFieldMapping(String avroFieldName, int avroPosition, Type icebergType, Schema avroType) { + if (Type.TypeID.TIMESTAMP.equals(icebergType.typeId()) + || Type.TypeID.TIME.equals(icebergType.typeId()) + || Type.TypeID.MAP.equals(icebergType.typeId()) + || Type.TypeID.LIST.equals(icebergType.typeId()) + || Type.TypeID.STRUCT.equals(icebergType.typeId())) { + avroType = resolveUnionElement(avroType); + } + return new FieldMapping(avroPosition, avroFieldName, icebergType, avroType); + } + + private Schema resolveUnionElement(Schema schema) { + if (schema.getType() != Schema.Type.UNION) { + return schema; + } + + // Collect all non-NULL types + List nonNullTypes = new ArrayList<>(); + for (Schema s : schema.getTypes()) { + if (s.getType() != NULL) { + nonNullTypes.add(s); + } + } + + if (nonNullTypes.isEmpty()) { + throw new IllegalArgumentException("UNION schema contains only NULL type: " + schema); + } else if (nonNullTypes.size() == 1) { + // Only unwrap UNION if it contains exactly one non-NULL type (optional union) + return nonNullTypes.get(0); + } else { + // Multiple non-NULL types: non-optional union not supported + throw new UnsupportedOperationException( + "Non-optional UNION with multiple non-NULL types is not supported. " + + "Found " + nonNullTypes.size() + " non-NULL types in UNION: " + schema); + } + } + + + /** + * Pre-computes RecordBinders for nested STRUCT fields. + */ + private Map precomputeBindersMap(TypeAdapter typeAdapter) { + Map binders = new IdentityHashMap<>(); + + for (FieldMapping mapping : fieldMappings) { + if (mapping != null) { + precomputeBindersForType(mapping.icebergType(), mapping.avroSchema(), binders, typeAdapter); + } + } + return binders; + } + + /** + * Recursively precomputes binders for a given Iceberg type and its corresponding Avro schema. + */ + private void precomputeBindersForType(Type icebergType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + if (icebergType.isPrimitiveType()) { + return; // No binders needed for primitive types + } + + if (icebergType.isStructType() && !avroSchema.isUnion()) { + createStructBinder(icebergType.asStructType(), avroSchema, binders, typeAdapter); + } else if (icebergType.isStructType() && avroSchema.isUnion()) { + createUnionStructBinders(icebergType.asStructType(), avroSchema, binders, typeAdapter); + } else if (icebergType.isListType()) { + createListBinder(icebergType.asListType(), avroSchema, binders, typeAdapter); + } else if (icebergType.isMapType()) { + createMapBinder(icebergType.asMapType(), avroSchema, binders, typeAdapter); + } + } + + /** + * Creates binders for STRUCT types represented as Avro UNIONs. + */ + private void createUnionStructBinders(Types.StructType structType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + org.apache.iceberg.Schema schema = structType.asSchema(); + SchemaBuilder.FieldAssembler schemaBuilder = SchemaBuilder.record(avroSchema.getName()).fields() + .name("tag").type().intType().noDefault(); + int tag = 0; + for (Schema unionMember : avroSchema.getTypes()) { + if (unionMember.getType() != NULL) { + schemaBuilder.name("field" + tag).type(unionMember).noDefault(); + tag++; + } + } + RecordBinder structBinder = new RecordBinder(schema, schemaBuilder.endRecord(), typeAdapter, batchFieldCount); + binders.put(avroSchema, structBinder); + } + + /** + * Creates a binder for a STRUCT type field. + */ + private void createStructBinder(Types.StructType structType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + org.apache.iceberg.Schema schema = structType.asSchema(); + RecordBinder structBinder = new RecordBinder(schema, avroSchema, typeAdapter, batchFieldCount); + binders.put(avroSchema, structBinder); + } + + /** + * Creates binders for LIST type elements (if they are STRUCT types). + */ + private void createListBinder(Types.ListType listType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + Type elementType = listType.elementType(); + if (elementType.isStructType()) { + Schema elementAvroSchema = avroSchema.getElementType(); + createStructBinder(elementType.asStructType(), elementAvroSchema, binders, typeAdapter); + } + } + + /** + * Creates binders for MAP type keys and values (if they are STRUCT types). + * Handles two Avro representations: ARRAY of key-value records, or native MAP. + */ + private void createMapBinder(Types.MapType mapType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + Type keyType = mapType.keyType(); + Type valueType = mapType.valueType(); + + if (ARRAY.equals(avroSchema.getType())) { + // Avro represents MAP as ARRAY of records with "key" and "value" fields + createMapAsArrayBinder(keyType, valueType, avroSchema, binders, typeAdapter); + } else { + // Avro represents MAP as native MAP type + createMapAsMapBinder(keyType, valueType, avroSchema, binders, typeAdapter); + } + } + + /** + * Handles MAP represented as Avro ARRAY of {key, value} records. + */ + private void createMapAsArrayBinder(Type keyType, Type valueType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + Schema elementSchema = avroSchema.getElementType(); + + // Process key if it's a STRUCT + if (keyType.isStructType()) { + Schema keyAvroSchema = elementSchema.getField("key").schema(); + createStructBinder(keyType.asStructType(), keyAvroSchema, binders, typeAdapter); + } + + // Process value if it's a STRUCT + if (valueType.isStructType()) { + Schema valueAvroSchema = elementSchema.getField("value").schema(); + createStructBinder(valueType.asStructType(), valueAvroSchema, binders, typeAdapter); + } + } + + /** + * Handles MAP represented as Avro native MAP type. + */ + private void createMapAsMapBinder(Type keyType, Type valueType, Schema avroSchema, + Map binders, + TypeAdapter typeAdapter) { + // Struct keys in native MAP are not supported by Avro + if (keyType.isStructType()) { + throw new UnsupportedOperationException("Struct keys in MAP types are not supported"); + } + + // Process value if it's a STRUCT + if (valueType.isStructType()) { + Schema valueAvroSchema = avroSchema.getValueType(); + createStructBinder(valueType.asStructType(), valueAvroSchema, binders, typeAdapter); + } + } + + private static class AvroRecordView implements Record { + private final GenericRecord avroRecord; + private final org.apache.iceberg.Schema icebergSchema; + private final TypeAdapter typeAdapter; + private final Map fieldNameToPosition; + private final FieldMapping[] fieldMappings; + private final Map nestedStructBinders; + private final RecordBinder parentBinder; + + AvroRecordView(GenericRecord avroRecord, + org.apache.iceberg.Schema icebergSchema, + TypeAdapter typeAdapter, + Map fieldNameToPosition, + FieldMapping[] fieldMappings, + Map nestedStructBinders, + RecordBinder parentBinder) { + this.avroRecord = avroRecord; + this.icebergSchema = icebergSchema; + this.typeAdapter = typeAdapter; + this.fieldNameToPosition = fieldNameToPosition; + this.fieldMappings = fieldMappings; + this.nestedStructBinders = nestedStructBinders; + this.parentBinder = parentBinder; + } + + @Override + public Object get(int pos) { + if (avroRecord == null) { + throw new IllegalStateException("Avro record is null"); + } + if (pos < 0 || pos >= fieldMappings.length) { + throw new IndexOutOfBoundsException("Field position " + pos + " out of bounds"); + } + + FieldMapping mapping = fieldMappings[pos]; + if (mapping == null) { + return null; + } + Object avroValue = avroRecord.get(mapping.avroPosition()); + if (avroValue == null) { + return null; + } + Object result = convert(avroValue, mapping.avroSchema(), mapping.icebergType()); + + // Calculate and accumulate field count + long fieldCount = calculateFieldCount(result, mapping.icebergType()); + parentBinder.addFieldCount(fieldCount); + + return result; + } + + public Object convert(Object sourceValue, Schema sourceSchema, Type targetType) { + if (targetType.typeId() == Type.TypeID.STRUCT) { + RecordBinder binder = nestedStructBinders.get(sourceSchema); + if (binder == null) { + throw new IllegalStateException("Missing nested binder for schema: " + sourceSchema); + } + return binder.bind((GenericRecord) sourceValue); + } + return typeAdapter.convert(sourceValue, (Schema) sourceSchema, targetType, this::convert); + } + + /** + * Calculates the field count for a converted value based on its size. + * Large fields are counted multiple times based on the size threshold. + */ + private long calculateFieldCount(Object value, Type icebergType) { + if (value == null) { + return 0; + } + + switch (icebergType.typeId()) { + case STRING: + return FieldMetric.count((CharSequence) value); + case BINARY: + return FieldMetric.count((ByteBuffer) value); + case FIXED: + return FieldMetric.count((byte[]) value); + case LIST: + return calculateListFieldCount(value, ((Types.ListType) icebergType).elementType()); + case MAP: + return calculateMapFieldCount(value, (Types.MapType) icebergType); + default: + return 1; // Struct or Primitive types count as 1 field + } + } + + /** + * Calculates field count for List values by summing element costs. + */ + private long calculateListFieldCount(Object list, Type elementType) { + if (list == null) { + return 0; + } + long total = 1; + if (list instanceof List) { + for (Object element : (List) list) { + total += calculateFieldCount(element, elementType); + } + } + return total; + } + + /** + * Calculates field count for Map values by summing key and value costs. + */ + private long calculateMapFieldCount(Object map, Types.MapType mapType) { + if (map == null) { + return 0; + } + + long total = 1; + if (map instanceof Map) { + Map typedMap = (Map) map; + if (typedMap.isEmpty()) { + return total; + } + for (Map.Entry entry : typedMap.entrySet()) { + total += calculateFieldCount(entry.getKey(), mapType.keyType()); + total += calculateFieldCount(entry.getValue(), mapType.valueType()); + } + } + return total; + } + + @Override + public Object getField(String name) { + Integer position = fieldNameToPosition.get(name); + return position != null ? get(position) : null; + } + + @Override + public Types.StructType struct() { + return icebergSchema.asStruct(); + } + + @Override + public int size() { + return icebergSchema.columns().size(); + } + + @Override + public T get(int pos, Class javaClass) { + return javaClass.cast(get(pos)); + } + + // Unsupported operations + @Override + public void setField(String name, Object value) { + throw new UnsupportedOperationException("Read-only"); + } + + @Override + public Record copy() { + throw new UnsupportedOperationException("Read-only"); + } + + @Override + public Record copy(Map overwriteValues) { + throw new UnsupportedOperationException("Read-only"); + } + + @Override + public void set(int pos, T value) { + throw new UnsupportedOperationException("Read-only"); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/binder/StructConverter.java b/core/src/main/java/kafka/automq/table/binder/StructConverter.java new file mode 100644 index 0000000000..87372cd33c --- /dev/null +++ b/core/src/main/java/kafka/automq/table/binder/StructConverter.java @@ -0,0 +1,27 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.binder; + +import org.apache.iceberg.types.Type; + +@FunctionalInterface +public interface StructConverter { + + Object convert(Object sourceValue, S sourceSchema, Type targetType); +} diff --git a/core/src/main/java/kafka/automq/table/binder/TypeAdapter.java b/core/src/main/java/kafka/automq/table/binder/TypeAdapter.java new file mode 100644 index 0000000000..2f5b06c670 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/binder/TypeAdapter.java @@ -0,0 +1,50 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.binder; + +import org.apache.iceberg.types.Type; + +/** + * Converts values between different schema systems. + * + * @param The source schema type (e.g., org.apache.avro.Schema) + */ +public interface TypeAdapter { + + /** + * Converts a source value to the target Iceberg type. + * + * @param sourceValue The source value + * @param sourceSchema The source schema + * @param targetType The target Iceberg type + * @return The converted value + */ + Object convert(Object sourceValue, S sourceSchema, Type targetType); + + /** + * Converts a source value to the target Iceberg type with support for recursive struct conversion. + * + * @param sourceValue The source value + * @param sourceSchema The source schema + * @param targetType The target Iceberg type + * @param structConverter A callback for converting nested STRUCT types + * @return The converted value + */ + Object convert(Object sourceValue, S sourceSchema, Type targetType, StructConverter structConverter); +} diff --git a/core/src/main/java/kafka/automq/table/coordinator/Checkpoint.java b/core/src/main/java/kafka/automq/table/coordinator/Checkpoint.java new file mode 100644 index 0000000000..f36848a62b --- /dev/null +++ b/core/src/main/java/kafka/automq/table/coordinator/Checkpoint.java @@ -0,0 +1,206 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.coordinator; + +import kafka.automq.table.events.AvroCodec; +import kafka.automq.table.events.Element; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +public class Checkpoint implements Element { + private Status status; + private UUID commitId; + private Long taskOffset; + private long[] nextOffsets; + private UUID lastCommitId; + private Long lastCommitTimestamp; + private long[] preCommitOffsets; + private final Schema avroSchema; + + private static final Schema AVRO_SCHEMA = SchemaBuilder.builder().record(Checkpoint.class.getName()) + .fields() + .name("status").type().intType().noDefault() + .name("commitId").type(UUID_SCHEMA).noDefault() + .name("taskOffset").type().longType().noDefault() + .name("nextOffsets").type().array().items().longType().noDefault() + .name("lastCommitId").type(UUID_SCHEMA).noDefault() + .name("lastCommitTimestamp").type().longType().noDefault() + .name("preCommitOffsets").type().array().items().longType().arrayDefault(Collections.emptyList()) + .endRecord(); + + public Checkpoint(Schema schema) { + this.avroSchema = schema; + } + + public Checkpoint(Status status, UUID commitId, Long taskOffset, long[] nextOffsets, + UUID lastCommitId, Long lastCommitTimestamp, long[] preCommitOffsets) { + this.status = status; + this.commitId = commitId; + this.taskOffset = taskOffset; + this.nextOffsets = nextOffsets; + this.lastCommitId = lastCommitId; + this.lastCommitTimestamp = lastCommitTimestamp; + this.preCommitOffsets = preCommitOffsets; + this.avroSchema = AVRO_SCHEMA; + } + + public Status status() { + return status; + } + + public UUID commitId() { + return commitId; + } + + public Long taskOffset() { + return taskOffset; + } + + public long[] nextOffsets() { + return nextOffsets; + } + + public UUID lastCommitId() { + return lastCommitId; + } + + public Long lastCommitTimestamp() { + return lastCommitTimestamp; + } + + public long[] preCommitOffsets() { + return preCommitOffsets; + } + + public static Checkpoint decode(ByteBuffer buf) { + try { + byte[] bytes = new byte[buf.remaining()]; + buf.get(bytes); + return AvroCodec.decode(bytes); + } catch (IOException e) { + throw new IllegalArgumentException(e); + } + } + + public byte[] encode() { + try { + return AvroCodec.encode(this); + } catch (IOException e) { + throw new IllegalArgumentException(e); + } + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: { + this.status = Status.fromCode((Integer) v); + break; + } + case 1: { + this.commitId = Element.toUuid((GenericData.Fixed) v); + break; + } + case 2: { + this.taskOffset = (Long) v; + break; + } + case 3: { + //noinspection unchecked + this.nextOffsets = ((List) v).stream().mapToLong(l -> l).toArray(); + break; + } + case 4: { + this.lastCommitId = Element.toUuid((GenericData.Fixed) v); + break; + } + case 5: { + this.lastCommitTimestamp = (Long) v; + break; + } + case 6: { + //noinspection unchecked + this.preCommitOffsets = ((List) v).stream().mapToLong(l -> l).toArray(); + break; + } + default: { + throw new IndexOutOfBoundsException("Invalid index: " + i); + } + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: { + return status.code(); + } + case 1: { + return Element.toFixed(commitId); + } + case 2: { + return taskOffset; + } + case 3: { + return Arrays.stream(nextOffsets).boxed().collect(Collectors.toList()); + } + case 4: { + return Element.toFixed(lastCommitId); + } + case 5: { + return lastCommitTimestamp; + } + case 6: { + return Arrays.stream(preCommitOffsets).boxed().collect(Collectors.toList()); + } + default: { + throw new IndexOutOfBoundsException("Invalid index: " + i); + } + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + @Override + public String toString() { + return "Checkpoint{" + + "status=" + status + + ", commitId=" + commitId + + ", taskOffset=" + taskOffset + + ", nextOffsets=" + Arrays.toString(nextOffsets) + + ", lastCommitId=" + lastCommitId + + ", lastCommitTimestamp=" + lastCommitTimestamp + + ", preCommitOffsets=" + Arrays.toString(preCommitOffsets) + + '}'; + } +} diff --git a/core/src/main/java/kafka/automq/table/coordinator/Status.java b/core/src/main/java/kafka/automq/table/coordinator/Status.java new file mode 100644 index 0000000000..6a2a019fb2 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/coordinator/Status.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.coordinator; + +public enum Status { + INIT(0), + REQUEST_COMMIT(1), + COMMITTED(2), + PRE_COMMIT(3); + + private final int code; + + Status(int code) { + this.code = code; + } + + public int code() { + return code; + } + + public static Status fromCode(int code) { + switch (code) { + case 0: + return INIT; + case 1: + return REQUEST_COMMIT; + case 2: + return COMMITTED; + case 3: + return PRE_COMMIT; + default: + throw new IllegalArgumentException("Invalid code: " + code); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/coordinator/TableCoordinator.java b/core/src/main/java/kafka/automq/table/coordinator/TableCoordinator.java new file mode 100644 index 0000000000..e56f755b8a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/coordinator/TableCoordinator.java @@ -0,0 +1,524 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.coordinator; + +import kafka.automq.table.Channel; +import kafka.automq.table.events.CommitRequest; +import kafka.automq.table.events.CommitResponse; +import kafka.automq.table.events.Envelope; +import kafka.automq.table.events.Errors; +import kafka.automq.table.events.Event; +import kafka.automq.table.events.EventType; +import kafka.automq.table.events.WorkerOffset; +import kafka.automq.table.metric.TableTopicMetricsManager; +import kafka.automq.table.utils.PartitionUtil; +import kafka.automq.table.utils.TableIdentifierUtil; +import kafka.log.streamaspect.MetaKeyValue; +import kafka.log.streamaspect.MetaStream; +import kafka.server.MetadataCache; + +import org.apache.kafka.storage.internals.log.LogConfig; + +import com.automq.stream.s3.metrics.Metrics; +import com.automq.stream.s3.metrics.TimerUtil; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.Time; +import com.automq.stream.utils.threads.EventLoop; + +import org.apache.commons.lang3.StringUtils; +import org.apache.iceberg.AppendFiles; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.RowDelta; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.Transaction; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.exceptions.NoSuchTableException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.Closeable; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.BitSet; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; + +@SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"}) +public class TableCoordinator implements Closeable { + private static final Logger LOGGER = LoggerFactory.getLogger(TableCoordinator.class); + private static final ScheduledExecutorService SCHEDULER = Threads.newSingleThreadScheduledExecutor("table-coordinator", true, LOGGER); + private static final ExecutorService EXPIRE_SNAPSHOT_EXECUTOR = Threads.newFixedThreadPoolWithMonitor(Systems.CPU_CORES * 2, "table-coordinator-expire-snapshot", true, LOGGER); + private static final String TABLE_COMMIT_CHECKPOINT_KEY = "TABLE_COMMIT_CHECKPOINT"; + private static final String SNAPSHOT_COMMIT_ID = "automq.commit.id"; + private static final String WATERMARK = "automq.watermark"; + private static final UUID NOOP_UUID = new UUID(0, 0); + private static final long NOOP_WATERMARK = -1L; + + private final Catalog catalog; + private final String topic; + private final String name; + private final MetaStream metaStream; + private final Channel channel; + private Table table; + private final EventLoop eventLoop; + private final MetadataCache metadataCache; + private final TableIdentifier tableIdentifier; + private final Time time = Time.SYSTEM; + private final long commitTimeout = TimeUnit.SECONDS.toMillis(30); + private volatile boolean closed = false; + private final Supplier config; + private final Metrics.LongGaugeBundle.LongGauge delayMetric; + private final Metrics.DoubleGaugeBundle.DoubleGauge fieldsPerSecondMetric; + + public TableCoordinator(Catalog catalog, String topic, MetaStream metaStream, Channel channel, + EventLoop eventLoop, MetadataCache metadataCache, Supplier config) { + this.catalog = catalog; + this.topic = topic; + this.name = topic; + this.metaStream = metaStream; + this.channel = channel; + this.eventLoop = eventLoop; + this.metadataCache = metadataCache; + this.config = config; + this.tableIdentifier = TableIdentifierUtil.of(config.get().tableTopicNamespace, topic); + this.delayMetric = TableTopicMetricsManager.registerDelay(topic); + this.fieldsPerSecondMetric = TableTopicMetricsManager.registerFieldsPerSecond(topic); + } + + private CommitStatusMachine commitStatusMachine; + + public void start() { + delayMetric.clear(); + fieldsPerSecondMetric.record(0.0); + + // await for a while to avoid multi coordinators concurrent commit. + SCHEDULER.schedule(() -> { + eventLoop.execute(() -> { + try { + // recover checkpoint from metaStream + Optional buf = metaStream.get(TABLE_COMMIT_CHECKPOINT_KEY); + if (buf.isPresent()) { + Checkpoint checkpoint = Checkpoint.decode(buf.get()); + commitStatusMachine = new CommitStatusMachine(checkpoint); + } else { + commitStatusMachine = new CommitStatusMachine(); + } + this.run(); + } catch (Throwable e) { + LOGGER.error("[TABLE_COORDINATOR_START_FAIL],{}", this, e); + } + }); + }, 10, TimeUnit.SECONDS); + } + + @Override + public void close() { + // quick close + closed = true; + delayMetric.close(); + fieldsPerSecondMetric.close(); + eventLoop.execute(() -> { + if (commitStatusMachine != null) { + commitStatusMachine.close(); + } + LOGGER.info("[TABLE_COORDINATOR_UNLOAD],{}", topic); + }); + } + + @Override + public String toString() { + return "TableCoordinator{" + + ", topic='" + topic + + '}'; + } + + private void run() { + eventLoop.execute(this::run0); + } + + private void run0() { + try { + if (closed) { + return; + } + switch (commitStatusMachine.status) { + case INIT: + case COMMITTED: + commitStatusMachine.nextRoundCommit(); + break; + case REQUEST_COMMIT: + commitStatusMachine.tryMoveToCommittedStatus(); + break; + default: + LOGGER.error("[TABLE_COORDINATOR_UNKNOWN_STATUS],{}", commitStatusMachine.status); + } + SCHEDULER.schedule(this::run, 1, TimeUnit.SECONDS); + } catch (Exception e) { + if (closed) { + LOGGER.warn("Error in table coordinator", e); + } else { + LOGGER.error("Error in table coordinator", e); + } + if (commitStatusMachine != null) { + commitStatusMachine.close(); + } + if (!closed) { + // reset the coordinator and retry later + SCHEDULER.schedule(this::start, 30, TimeUnit.SECONDS); + } + } + } + + static class CommitInfo { + UUID commitId; + long taskOffset; + long[] nextOffsets; + + public CommitInfo(UUID commitId, long taskOffset, long[] nextOffsets) { + this.commitId = commitId; + this.taskOffset = taskOffset; + this.nextOffsets = nextOffsets; + } + } + + class CommitStatusMachine { + Status status = Status.INIT; + + CommitInfo last; + CommitInfo processing; + + long lastCommitTimestamp = 0; + + List dataFiles = new ArrayList<>(); + List deleteFiles = new ArrayList<>(); + private BitSet readyPartitions; + private int unreadyPartitionCount; + private long requestCommitTimestamp = time.milliseconds(); + private long commitFieldCount; + private long[] partitionWatermarks = new long[0]; + private boolean fastNextCommit; + private long polledOffset = 0; + + String lastAppliedPartitionBy; + + Channel.SubChannel subChannel; + + public CommitStatusMachine() { + last = new CommitInfo(NOOP_UUID, 0, new long[0]); + this.subChannel = channel.subscribeData(topic, Channel.LATEST_OFFSET); + } + + public CommitStatusMachine(Checkpoint checkpoint) { + status = checkpoint.status(); + lastCommitTimestamp = checkpoint.lastCommitTimestamp(); + switch (status) { + case INIT: { + throw new IllegalStateException("Invalid checkpoint status: " + status); + } + case REQUEST_COMMIT: { + last = new CommitInfo(checkpoint.lastCommitId(), checkpoint.taskOffset(), copy(checkpoint.nextOffsets())); + processing = new CommitInfo(checkpoint.commitId(), checkpoint.taskOffset(), copy(checkpoint.nextOffsets())); + break; + } + case PRE_COMMIT: { + last = new CommitInfo(checkpoint.lastCommitId(), checkpoint.taskOffset(), copy(checkpoint.nextOffsets())); + processing = new CommitInfo(checkpoint.commitId(), checkpoint.taskOffset(), copy(checkpoint.nextOffsets())); + // check if the last commit is the same as the last snapshot commit + Optional lastSnapshotCommitId = getLastSnapshotCommitId(); + if (lastSnapshotCommitId.isPresent() && Objects.equals(processing.commitId, lastSnapshotCommitId.get())) { + // TODO: for true exactly once, iceberg should support conditional commit + LOGGER.info("[ALREADY_COMMITED],{}", processing.commitId); + last = new CommitInfo(checkpoint.commitId(), checkpoint.taskOffset(), copy(checkpoint.preCommitOffsets())); + processing = null; + status = Status.COMMITTED; + } else { + status = Status.REQUEST_COMMIT; + } + break; + } + case COMMITTED: { + last = new CommitInfo(checkpoint.lastCommitId(), checkpoint.taskOffset(), copy(checkpoint.nextOffsets())); + processing = null; + break; + } + } + + if (status == Status.REQUEST_COMMIT) { + //noinspection DataFlowIssue + int partitionCount = processing.nextOffsets.length; + readyPartitions = new BitSet(partitionCount); + unreadyPartitionCount = partitionCount; + partitionWatermarks = new long[partitionCount]; + Arrays.fill(partitionWatermarks, NOOP_WATERMARK); + } + this.subChannel = channel.subscribeData(topic, Optional.ofNullable(processing).map(p -> p.taskOffset).orElse(last.taskOffset)); + LOGGER.info("[LOAD_CHECKPOINT],{},{}", topic, checkpoint); + } + + public void nextRoundCommit() throws Exception { + // drain the sub channel to avoid catchup-read + for (; ; ) { + Envelope envelope = subChannel.poll(); + if (envelope == null) { + break; + } + polledOffset = envelope.offset(); + } + long nextCheck = config.get().tableTopicCommitInterval - (time.milliseconds() - lastCommitTimestamp); + if (!fastNextCommit && nextCheck > 0) { + return; + } + last.taskOffset = polledOffset; + processing = new CommitInfo(UUID.randomUUID(), polledOffset, copy(last.nextOffsets)); + int partitionNum = (Integer) metadataCache.numPartitions(topic).get(); + handlePartitionNumChange(partitionNum); + List workerOffsets = new ArrayList<>(partitionNum); + for (int partitionId = 0; partitionId < partitionNum; partitionId++) { + int leaderEpoch = metadataCache.getPartitionInfo(topic, partitionId).get().leaderEpoch(); + long offset = processing.nextOffsets[partitionId]; + workerOffsets.add(new WorkerOffset(partitionId, leaderEpoch, offset)); + } + + status = Status.REQUEST_COMMIT; + dataFiles = new ArrayList<>(); + deleteFiles = new ArrayList<>(); + readyPartitions = new BitSet(partitionNum); + unreadyPartitionCount = partitionNum; + requestCommitTimestamp = time.milliseconds(); + commitFieldCount = 0; + fastNextCommit = false; + + Checkpoint checkpoint = new Checkpoint(status, processing.commitId, polledOffset, + processing.nextOffsets, last.commitId, lastCommitTimestamp, new long[0]); + metaStream.append(MetaKeyValue.of(TABLE_COMMIT_CHECKPOINT_KEY, ByteBuffer.wrap(checkpoint.encode()))).get(); + int specId = Optional.ofNullable(table).map(t -> t.spec().specId()).orElse(CommitRequest.NOOP_SPEC_ID); + CommitRequest commitRequest = new CommitRequest(processing.commitId, topic, specId, workerOffsets); + LOGGER.info("[SEND_COMMIT_REQUEST],{},{}", name, commitRequest); + channel.send(topic, new Event(time.milliseconds(), EventType.COMMIT_REQUEST, commitRequest)); + } + + public void tryMoveToCommittedStatus() throws Exception { + for (; ; ) { + boolean awaitCommitTimeout = (time.milliseconds() - requestCommitTimestamp) > commitTimeout; + if (!awaitCommitTimeout) { + for (; ; ) { + Envelope envelope = subChannel.poll(); + if (envelope == null) { + break; + } else { + polledOffset = envelope.offset(); + CommitResponse commitResponse = envelope.event().payload(); + if (!processing.commitId.equals(commitResponse.commitId())) { + continue; + } + LOGGER.info("[RECEIVE_COMMIT_RESPONSE],{}", commitResponse); + dataFiles.addAll(commitResponse.dataFiles()); + deleteFiles.addAll(commitResponse.deleteFiles()); + boolean moveNextOffset = commitResponse.code() == Errors.NONE || commitResponse.code() == Errors.MORE_DATA; + for (WorkerOffset nextOffset : commitResponse.nextOffsets()) { + if (moveNextOffset) { + processing.nextOffsets[nextOffset.partition()] = Math.max(nextOffset.offset(), processing.nextOffsets[nextOffset.partition()]); + } + if (!readyPartitions.get(nextOffset.partition())) { + readyPartitions.set(nextOffset.partition()); + unreadyPartitionCount--; + } + } + + commitFieldCount += commitResponse.topicMetric().fieldCount(); + commitResponse.partitionMetrics().forEach(m -> { + if (m.watermark() != NOOP_WATERMARK) { + partitionWatermarks[m.partition()] = Math.max(partitionWatermarks[m.partition()], m.watermark()); + } + }); + + if (commitResponse.code() == Errors.MORE_DATA) { + fastNextCommit = true; + } + } + } + if (unreadyPartitionCount != 0) { + break; + } + } + TimerUtil transactionCommitTimer = new TimerUtil(); + Transaction transaction = null; + if (!dataFiles.isEmpty() || !deleteFiles.isEmpty()) { + transaction = getTable().newTransaction(); + long watermark = watermark(partitionWatermarks); + if (deleteFiles.isEmpty()) { + AppendFiles appendFiles = transaction.newAppend(); + appendFiles.set(SNAPSHOT_COMMIT_ID, processing.commitId.toString()); + appendFiles.set(WATERMARK, Long.toString(watermark)); + dataFiles.forEach(appendFiles::appendFile); + appendFiles.commit(); + } else { + RowDelta delta = transaction.newRowDelta(); + delta.set(SNAPSHOT_COMMIT_ID, processing.commitId.toString()); + delta.set(WATERMARK, Long.toString(watermark)); + dataFiles.forEach(delta::addRows); + deleteFiles.forEach(delta::addDeletes); + delta.commit(); + } + try { + LogConfig currentLogConfig = config.get(); + if (currentLogConfig.tableTopicExpireSnapshotEnabled) { + transaction.expireSnapshots() + .expireOlderThan(System.currentTimeMillis() - TimeUnit.HOURS.toMillis(currentLogConfig.tableTopicExpireSnapshotOlderThanHours)) + .retainLast(currentLogConfig.tableTopicExpireSnapshotRetainLast) + .executeDeleteWith(EXPIRE_SNAPSHOT_EXECUTOR) + .commit(); + } + } catch (Exception exception) { + // skip expire snapshot failure + LOGGER.error("[EXPIRE_SNAPSHOT_FAIL],{}", getTable().name(), exception); + } + } + + recordMetrics(); + if (transaction != null) { + Checkpoint checkpoint = new Checkpoint(Status.PRE_COMMIT, processing.commitId, last.taskOffset, + last.nextOffsets, last.commitId, lastCommitTimestamp, processing.nextOffsets); + metaStream.append(MetaKeyValue.of(TABLE_COMMIT_CHECKPOINT_KEY, ByteBuffer.wrap(checkpoint.encode()))).get(); + transaction.commitTransaction(); + if (awaitCommitTimeout) { + LOGGER.warn("[COMMIT_AWAIT_TIMEOUT],{}", processing.commitId); + } + } + lastCommitTimestamp = time.milliseconds(); + + status = Status.COMMITTED; + Checkpoint checkpoint = new Checkpoint(status, NOOP_UUID, polledOffset, + processing.nextOffsets, processing.commitId, lastCommitTimestamp, new long[0]); + metaStream.append(MetaKeyValue.of(TABLE_COMMIT_CHECKPOINT_KEY, ByteBuffer.wrap(checkpoint.encode()))).get(); + LOGGER.info("[COMMIT_COMPLETE],{},{},commitCost={}ms,total={}ms", name, processing.commitId, + transactionCommitTimer.elapsedAs(TimeUnit.MILLISECONDS), lastCommitTimestamp - requestCommitTimestamp); + last = processing; + processing = null; + + if (awaitCommitTimeout) { + fastNextCommit = true; + } + + if (tryEvolvePartition()) { + // Let workers know the partition evolution. + fastNextCommit = true; + LOGGER.info("[TABLE_PARTITION_EVOLUTION],{}", config.get().tableTopicPartitionBy); + } + + break; + } + } + + public void close() { + if (subChannel != null) { + subChannel.close(); + } + } + + private Table getTable() { + if (table == null) { + table = catalog.loadTable(tableIdentifier); + } + return table; + } + + private void handlePartitionNumChange(int partitionNum) { + if (partitionNum > processing.nextOffsets.length) { + long[] newNextOffsets = new long[partitionNum]; + System.arraycopy(processing.nextOffsets, 0, newNextOffsets, 0, processing.nextOffsets.length); + processing.nextOffsets = newNextOffsets; + } + if (partitionNum > partitionWatermarks.length) { + long[] newPartitionWatermarks = new long[partitionNum]; + System.arraycopy(partitionWatermarks, 0, newPartitionWatermarks, 0, partitionWatermarks.length); + for (int i = partitionWatermarks.length; i < partitionNum; i++) { + newPartitionWatermarks[i] = NOOP_WATERMARK; + } + partitionWatermarks = newPartitionWatermarks; + } + } + + private void recordMetrics() { + long now = System.currentTimeMillis(); + double fps = commitFieldCount * 1000.0 / Math.max(now - lastCommitTimestamp, 1); + fieldsPerSecondMetric.record(fps); + long watermarkTimestamp = watermark(partitionWatermarks); + if (watermarkTimestamp == NOOP_WATERMARK) { + delayMetric.clear(); + } else { + delayMetric.record(Math.max(now - watermarkTimestamp, 0)); + } + } + + private boolean tryEvolvePartition() { + String newPartitionBy = config.get().tableTopicPartitionBy; + if (Objects.equals(newPartitionBy, lastAppliedPartitionBy) || StringUtils.isBlank(newPartitionBy) || table == null) { + return false; + } + boolean changed = PartitionUtil.evolve(PartitionUtil.parsePartitionBy(newPartitionBy), table); + lastAppliedPartitionBy = newPartitionBy; + return changed; + } + + private Optional getLastSnapshotCommitId() { + try { + if (table == null) { + table = catalog.loadTable(tableIdentifier); + } + Snapshot snapshot = table.currentSnapshot(); + return Optional.ofNullable(snapshot).map(s -> s.summary().get(SNAPSHOT_COMMIT_ID)).map(UUID::fromString); + } catch (NoSuchTableException ex) { + return Optional.empty(); + } + } + + } + + static long watermark(long[] watermarks) { + boolean match = false; + long watermark = Long.MAX_VALUE; + for (long w : watermarks) { + if (w != NOOP_WATERMARK) { + match = true; + watermark = Math.min(watermark, w); + } + } + if (!match) { + return NOOP_WATERMARK; + } + return watermark; + } + + static long[] copy(long[] array) { + return Arrays.copyOf(array, array.length); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/SchemaResolutionResolver.java b/core/src/main/java/kafka/automq/table/deserializer/SchemaResolutionResolver.java new file mode 100644 index 0000000000..e95e3e1f60 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/SchemaResolutionResolver.java @@ -0,0 +1,93 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer; + +import kafka.automq.table.deserializer.proto.schema.MessageIndexes; + +import java.nio.ByteBuffer; + + +/** + * Interface for resolving schema information for protobuf message deserialization. + * This interface supports different strategies for obtaining schema ID and message structure: + * - Parse from message header (standard Confluent format) + * - Lookup latest schema from registry by subject name + */ +public interface SchemaResolutionResolver { + + /** + * Resolves schema information for the given message payload and topic. + * + * @param topic The Kafka topic name + * @param payload The serialized protobuf payload as ByteBuffer + * @return SchemaResolution containing schema ID, message indexes, and message bytes buffer + * @throws org.apache.kafka.common.errors.SerializationException if resolution fails + */ + SchemaResolution resolve(String topic, ByteBuffer payload); + + + int getSchemaId(String topic, ByteBuffer payload); + + /** + * Container class for resolved schema information. + */ + class SchemaResolution { + private final int schemaId; + private final ByteBuffer messageBytes; + + private MessageIndexes indexes; + private String subject; + private String messageTypeName; + + public SchemaResolution(int schemaId, MessageIndexes indexes, ByteBuffer messageBytes) { + this.schemaId = schemaId; + this.indexes = indexes; + this.messageBytes = messageBytes; + } + + public SchemaResolution(int schemaId, String subject, String messageTypeName, ByteBuffer messageBytes) { + this.schemaId = schemaId; + this.subject = subject; + this.messageTypeName = messageTypeName; + this.messageBytes = messageBytes; + } + + public int getSchemaId() { + return schemaId; + } + + public MessageIndexes getIndexes() { + return indexes; + } + + + public ByteBuffer getMessageBytes() { + return messageBytes; + } + + public String getMessageTypeName() { + return messageTypeName; + } + + public String getSubject() { + return subject; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/AbstractCustomKafkaProtobufDeserializer.java b/core/src/main/java/kafka/automq/table/deserializer/proto/AbstractCustomKafkaProtobufDeserializer.java new file mode 100644 index 0000000000..42898421ee --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/AbstractCustomKafkaProtobufDeserializer.java @@ -0,0 +1,188 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import kafka.automq.table.deserializer.SchemaResolutionResolver; + +import org.apache.kafka.common.errors.InvalidConfigurationException; +import org.apache.kafka.common.errors.SerializationException; +import org.apache.kafka.common.errors.TimeoutException; +import org.apache.kafka.common.header.Headers; + +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Message; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InterruptedIOException; +import java.nio.ByteBuffer; +import java.util.Map; +import java.util.Objects; + +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; +import io.confluent.kafka.schemaregistry.utils.BoundedConcurrentHashMap; +import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDe; + +public abstract class AbstractCustomKafkaProtobufDeserializer + extends AbstractKafkaSchemaSerDe { + + protected final Map schemaCache; + protected final SchemaResolutionResolver schemaResolutionResolver; + + public AbstractCustomKafkaProtobufDeserializer() { + this.schemaCache = new BoundedConcurrentHashMap<>(1000); + this.schemaResolutionResolver = new HeaderBasedSchemaResolutionResolver(); + } + + public AbstractCustomKafkaProtobufDeserializer(SchemaResolutionResolver schemaResolutionResolver) { + this.schemaCache = new BoundedConcurrentHashMap<>(1000); + this.schemaResolutionResolver = schemaResolutionResolver; + } + + protected void configure(CustomKafkaProtobufDeserializerConfig config) { + configureClientProperties(config, new ProtobufSchemaProvider()); + } + + /** + * Deserialize protobuf message from the given byte array. + * The implementation follows the open-closed principle by breaking down the + * deserialization process into multiple phases that can be extended by subclasses. + * + * @param topic The Kafka topic + * @param headers The Kafka record headers + * @param payload The serialized protobuf payload + * @return The deserialized object + */ + protected T deserialize(String topic, Headers headers, byte[] payload) + throws SerializationException, InvalidConfigurationException { + // Phase 1: Pre-validation + if (payload == null) { + return null; + } + if (schemaRegistry == null) { + throw new InvalidConfigurationException("Schema registry not found, make sure the schema.registry.url is set"); + } + + try { + // Phase 2: Schema Resolution + ByteBuffer byteBuffer = ByteBuffer.wrap(payload); + SchemaResolutionResolver.SchemaResolution resolution = schemaResolutionResolver.resolve(topic, byteBuffer); + int schemaId = resolution.getSchemaId(); + ByteBuffer messageBytes = resolution.getMessageBytes(); + + // Phase 3: Schema Processing + ProtobufSchemaWrapper protobufSchemaWrapper = processSchema(topic, schemaId, resolution); + Descriptors.Descriptor targetDescriptor = protobufSchemaWrapper.getDescriptor(); + + // Phase 4: Message Deserialization + Message message = deserializeMessage(targetDescriptor, messageBytes); + + @SuppressWarnings("unchecked") + T result = (T) message; + return result; + } catch (InterruptedIOException e) { + throw new TimeoutException("Error deserializing Protobuf message", e); + } catch (IOException | RuntimeException e) { + throw new SerializationException("Error deserializing Protobuf message", e); + } + } + + private Message deserializeMessage(Descriptors.Descriptor descriptor, ByteBuffer messageBytes) throws IOException { + if (descriptor == null) { + throw new SerializationException("No Protobuf Descriptor found"); + } + + // Convert ByteBuffer to byte array for DynamicMessage.parseFrom + byte[] bytes; + if (messageBytes.hasArray() && messageBytes.arrayOffset() == 0 && messageBytes.remaining() == messageBytes.array().length) { + // Use the backing array directly if it's a simple case + bytes = messageBytes.array(); + } else { + // Create a new byte array for the remaining bytes + bytes = new byte[messageBytes.remaining()]; + messageBytes.duplicate().get(bytes); + } + + return DynamicMessage.parseFrom(descriptor, new ByteArrayInputStream(bytes)); + } + + /** + * Phase 3: Process and retrieve the schema + */ + protected ProtobufSchemaWrapper processSchema(String topic, int schemaId, SchemaResolutionResolver.SchemaResolution resolution) { + String subject = resolution.getSubject() == null ? + getSubjectName(topic, isKey, null, null) : resolution.getSubject(); + SchemaKey key = new SchemaKey(subject, schemaId); + try { + CustomProtobufSchema schema = (CustomProtobufSchema) schemaRegistry.getSchemaBySubjectAndId(subject, schemaId); + return schemaCache.computeIfAbsent(key, k -> { + if (resolution.getIndexes() != null) { + return new ProtobufSchemaWrapper(schema, resolution.getIndexes()); + } else { + return new ProtobufSchemaWrapper(schema, resolution.getMessageTypeName()); + } + }); + } catch (IOException | RestClientException e) { + throw new SerializationException("Error retrieving Protobuf schema for id " + schemaId, e); + } + } + protected static final class SchemaKey { + private final String subject; + private final int schemaId; + + protected SchemaKey(String subject, int schemaId) { + this.subject = subject; + this.schemaId = schemaId; + } + + public String subject() { + return subject; + } + + public int schemaId() { + return schemaId; + } + + @Override + public boolean equals(Object obj) { + if (obj == this) + return true; + if (obj == null || obj.getClass() != this.getClass()) + return false; + var that = (SchemaKey) obj; + return Objects.equals(this.subject, that.subject) && + this.schemaId == that.schemaId; + } + + @Override + public int hashCode() { + return Objects.hash(subject, schemaId); + } + + @Override + public String toString() { + return "SchemaKey[" + + "subject=" + subject + ", " + + "schemaId=" + schemaId + ']'; + } + + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/CustomKafkaProtobufDeserializer.java b/core/src/main/java/kafka/automq/table/deserializer/proto/CustomKafkaProtobufDeserializer.java new file mode 100644 index 0000000000..04f0db7d3f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/CustomKafkaProtobufDeserializer.java @@ -0,0 +1,77 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import kafka.automq.table.deserializer.SchemaResolutionResolver; + +import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.serialization.Deserializer; + +import com.google.protobuf.Message; + +import java.io.IOException; +import java.util.Map; + +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; + +public class CustomKafkaProtobufDeserializer + extends AbstractCustomKafkaProtobufDeserializer implements Deserializer { + + public CustomKafkaProtobufDeserializer() { + } + + public CustomKafkaProtobufDeserializer(SchemaResolutionResolver resolver) { + super(resolver); + } + + public CustomKafkaProtobufDeserializer(SchemaRegistryClient schemaRegistry) { + this.schemaRegistry = schemaRegistry; + } + + public CustomKafkaProtobufDeserializer(SchemaRegistryClient schemaRegistry, SchemaResolutionResolver schemaResolutionResolver) { + super(schemaResolutionResolver); + this.schemaRegistry = schemaRegistry; + } + + @Override + public void configure(Map configs, boolean isKey) { + CustomKafkaProtobufDeserializerConfig config = new CustomKafkaProtobufDeserializerConfig(configs); + configure(config); + } + + @Override + public T deserialize(String topic, byte[] bytes) { + return this.deserialize(topic, null, bytes); + } + + @Override + public T deserialize(String topic, Headers headers, byte[] data) { + return (T) super.deserialize(topic, headers, data); + } + + @Override + public void close() { + try { + super.close(); + } catch (IOException e) { + throw new RuntimeException("Exception while closing deserializer", e); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/CustomKafkaProtobufDeserializerConfig.java b/core/src/main/java/kafka/automq/table/deserializer/proto/CustomKafkaProtobufDeserializerConfig.java new file mode 100644 index 0000000000..a87669376f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/CustomKafkaProtobufDeserializerConfig.java @@ -0,0 +1,33 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import org.apache.kafka.common.config.ConfigDef; + +import java.util.Map; + +import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig; + +public class CustomKafkaProtobufDeserializerConfig extends AbstractKafkaSchemaSerDeConfig { + private static final ConfigDef CONFIG_DEF = baseConfigDef(); + public CustomKafkaProtobufDeserializerConfig(Map props) { + super(CONFIG_DEF, props); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/CustomProtobufSchema.java b/core/src/main/java/kafka/automq/table/deserializer/proto/CustomProtobufSchema.java new file mode 100644 index 0000000000..29b78c5093 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/CustomProtobufSchema.java @@ -0,0 +1,122 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import java.util.List; +import java.util.Map; +import java.util.Set; + +import io.confluent.kafka.schemaregistry.ParsedSchema; +import io.confluent.kafka.schemaregistry.client.rest.entities.Metadata; +import io.confluent.kafka.schemaregistry.client.rest.entities.RuleSet; +import io.confluent.kafka.schemaregistry.client.rest.entities.SchemaEntity; +import io.confluent.kafka.schemaregistry.client.rest.entities.SchemaReference; + +public class CustomProtobufSchema implements ParsedSchema { + private final String name; + private final Integer version; + private final Metadata metadata; + private final RuleSet ruleSet; + private final String schemaDefinition; + private final List references; + private final Map resolvedReferences; + + public CustomProtobufSchema( + String name, Integer version, + Metadata metadata, RuleSet ruleSet, + String schemaDefinition, List references, Map resolvedReferences) { + this.name = name; + this.version = version; + this.metadata = metadata; + this.ruleSet = ruleSet; + this.schemaDefinition = schemaDefinition; + this.references = references; + this.resolvedReferences = resolvedReferences; + } + + @Override + public String schemaType() { + return "PROTOBUF"; + } + + @Override + public String name() { + return name; + } + + @Override + public String canonicalString() { + return schemaDefinition; + } + + @Override + public Integer version() { + return version; + } + + @Override + public List references() { + return references; + } + + public Map resolvedReferences() { + return resolvedReferences; + } + + @Override + public Metadata metadata() { + return metadata; + } + + @Override + public RuleSet ruleSet() { + return ruleSet; + } + + @Override + public ParsedSchema copy() { + return new CustomProtobufSchema(name, version, metadata, ruleSet, schemaDefinition, references, resolvedReferences); + } + + @Override + public ParsedSchema copy(Integer version) { + return new CustomProtobufSchema(name, version, metadata, ruleSet, schemaDefinition, references, resolvedReferences); + } + + @Override + public ParsedSchema copy(Metadata metadata, RuleSet ruleSet) { + return new CustomProtobufSchema(name, version, metadata, ruleSet, schemaDefinition, references, resolvedReferences); + } + + @Override + public ParsedSchema copy(Map> tagsToAdd, Map> tagsToRemove) { + throw new UnsupportedOperationException("Not implemented"); + } + + @Override + public List isBackwardCompatible(ParsedSchema previousSchema) { + throw new UnsupportedOperationException("Not implemented"); + } + + @Override + public Object rawSchema() { + throw new UnsupportedOperationException("Not implemented"); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/HeaderBasedSchemaResolutionResolver.java b/core/src/main/java/kafka/automq/table/deserializer/proto/HeaderBasedSchemaResolutionResolver.java new file mode 100644 index 0000000000..0fb803975a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/HeaderBasedSchemaResolutionResolver.java @@ -0,0 +1,80 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import kafka.automq.table.deserializer.SchemaResolutionResolver; +import kafka.automq.table.deserializer.proto.schema.MessageIndexes; +import kafka.automq.table.process.exception.InvalidDataException; + +import org.apache.kafka.common.errors.SerializationException; + +import java.nio.ByteBuffer; + +/** + * Default implementation of SchemaResolutionResolver that parses schema information from message headers. + * This implementation handles the standard Confluent Kafka protobuf message format with magic byte, + * schema ID, message indexes, and message payload. + */ +public class HeaderBasedSchemaResolutionResolver implements SchemaResolutionResolver { + + private static final int SCHEMA_ID_SIZE = 4; + private static final int HEADER_SIZE = SCHEMA_ID_SIZE + 1; // magic byte + schema id + private static final byte MAGIC_BYTE = 0x0; + + @Override + public SchemaResolution resolve(String topic, ByteBuffer payload) { + if (payload == null) { + throw new SerializationException("Payload cannot be null"); + } + + ByteBuffer buffer = payload.duplicate(); + + int schemaId = readSchemaId(buffer); + + // Extract message indexes + MessageIndexes indexes = MessageIndexes.readFrom(buffer); + + // Extract message bytes as a slice of the buffer + ByteBuffer messageBytes = buffer.slice(); + + return new SchemaResolution(schemaId, indexes, messageBytes); + } + + @Override + public int getSchemaId(String topic, ByteBuffer payload) { + // io.confluent.kafka.serializers.DeserializationContext#constructor + return readSchemaId(payload.duplicate()); + } + + private int readSchemaId(ByteBuffer buffer) { + if (buffer.remaining() < HEADER_SIZE) { + throw new InvalidDataException("Invalid payload size: " + buffer.remaining() + ", expected at least " + HEADER_SIZE); + } + + // Extract magic byte + byte magicByte = buffer.get(); + if (magicByte != MAGIC_BYTE) { + throw new InvalidDataException("Unknown magic byte: " + magicByte); + } + + // Extract schema ID + return buffer.getInt(); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/LatestSchemaResolutionResolver.java b/core/src/main/java/kafka/automq/table/deserializer/proto/LatestSchemaResolutionResolver.java new file mode 100644 index 0000000000..95db40b34a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/LatestSchemaResolutionResolver.java @@ -0,0 +1,125 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import kafka.automq.table.deserializer.SchemaResolutionResolver; +import kafka.automq.table.deserializer.proto.schema.MessageIndexes; + +import org.apache.kafka.common.errors.SerializationException; + +import com.automq.stream.utils.Time; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collections; + +import io.confluent.kafka.schemaregistry.client.SchemaMetadata; +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; + +/** + * Implementation of SchemaResolutionResolver that retrieves the latest schema from Schema Registry by subject name. + * This implementation includes caching mechanism to avoid frequent registry queries. + * Cache entries are refreshed every 5 minutes. + */ +public class LatestSchemaResolutionResolver implements SchemaResolutionResolver { + private static final Logger log = LoggerFactory.getLogger(LatestSchemaResolutionResolver.class); + + private static final long CACHE_REFRESH_INTERVAL_MS = 5 * 60 * 1000; // 5 minutes + private static final MessageIndexes DEFAULT_INDEXES = new MessageIndexes(Collections.singletonList(0)); + + private final SchemaRegistryClient schemaRegistry; + + private volatile LatestSchemaResolutionResolver.CachedSchemaInfo schemaCache; + + private final Time time; + private final String subject; + private final String messageFullName; + + public LatestSchemaResolutionResolver(SchemaRegistryClient schemaRegistry, String subject) { + this(schemaRegistry, subject, null); + } + + public LatestSchemaResolutionResolver(SchemaRegistryClient schemaRegistry, String subject, String messageFullName) { + this.schemaRegistry = schemaRegistry; + this.time = Time.SYSTEM; + this.subject = subject; + this.messageFullName = messageFullName; + } + + @Override + public SchemaResolution resolve(String topic, ByteBuffer payload) { + if (payload == null) { + throw new SerializationException("Payload cannot be null"); + } + LatestSchemaResolutionResolver.CachedSchemaInfo cachedInfo = getCachedSchemaInfo(subject); + + if (messageFullName == null) { + return new SchemaResolution(cachedInfo.schemaId, DEFAULT_INDEXES, payload); + } else { + return new SchemaResolution(cachedInfo.schemaId, subject, messageFullName, payload); + } + } + + @Override + public int getSchemaId(String topic, ByteBuffer payload) { + LatestSchemaResolutionResolver.CachedSchemaInfo cachedInfo = getCachedSchemaInfo(subject); + return cachedInfo.schemaId; + } + + private LatestSchemaResolutionResolver.CachedSchemaInfo getCachedSchemaInfo(String subject) { + long currentTime = time.milliseconds(); + // First check (no lock) + if (schemaCache == null || currentTime - schemaCache.lastUpdated > CACHE_REFRESH_INTERVAL_MS) { + synchronized (this) { + // Second check (with lock) + if (schemaCache == null || currentTime - schemaCache.lastUpdated > CACHE_REFRESH_INTERVAL_MS) { + try { + SchemaMetadata latestSchema = schemaRegistry.getLatestSchemaMetadata(subject); + schemaCache = new LatestSchemaResolutionResolver.CachedSchemaInfo(latestSchema.getId(), currentTime); + } catch (IOException | RestClientException e) { + if (schemaCache == null) { + // No cached data and fresh fetch failed - this is a hard error + throw new SerializationException("Error retrieving schema for subject " + subject + + " and no cached data available", e); + } else { + log.warn("Failed to retrieve latest schema for subject '{}'. Using stale cache.", subject, e); + } + } + } + } + } + return schemaCache; + } + + + private static class CachedSchemaInfo { + final int schemaId; + final long lastUpdated; + + CachedSchemaInfo(int schemaId, long lastUpdated) { + this.schemaId = schemaId; + this.lastUpdated = lastUpdated; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/ProtobufSchemaProvider.java b/core/src/main/java/kafka/automq/table/deserializer/proto/ProtobufSchemaProvider.java new file mode 100644 index 0000000000..a8c34fbbd1 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/ProtobufSchemaProvider.java @@ -0,0 +1,48 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import java.util.Map; + +import io.confluent.kafka.schemaregistry.AbstractSchemaProvider; +import io.confluent.kafka.schemaregistry.ParsedSchema; +import io.confluent.kafka.schemaregistry.client.rest.entities.Schema; + +public class ProtobufSchemaProvider extends AbstractSchemaProvider { + + @Override + public String schemaType() { + return "PROTOBUF"; + } + + @Override + public ParsedSchema parseSchemaOrElseThrow(Schema schema, boolean isNew, boolean normalize) { + Map resolveReferences = resolveReferences(schema); + return new CustomProtobufSchema( + null, + schema.getVersion(), + schema.getMetadata(), + schema.getRuleSet(), + schema.getSchema(), + schema.getReferences(), + resolveReferences + ); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/ProtobufSchemaWrapper.java b/core/src/main/java/kafka/automq/table/deserializer/proto/ProtobufSchemaWrapper.java new file mode 100644 index 0000000000..8260466687 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/ProtobufSchemaWrapper.java @@ -0,0 +1,168 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto; + +import kafka.automq.table.deserializer.proto.parse.ProtobufSchemaParser; +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; +import kafka.automq.table.deserializer.proto.schema.MessageIndexes; + +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; +import com.google.protobuf.InvalidProtocolBufferException; +import com.squareup.wire.schema.internal.parser.MessageElement; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.ProtoParser; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.Base64; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.logging.Level; +import java.util.logging.Logger; + +import static kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants.DEFAULT_LOCATION; + +public class ProtobufSchemaWrapper { + private static final Logger LOGGER = Logger.getLogger(ProtobufSchemaWrapper.class.getName()); + + private final CustomProtobufSchema schema; + + private MessageIndexes messageIndexes; + private String messageTypeName; + + private static final Base64.Decoder DECODER = Base64.getDecoder(); + private ProtoFileElement rootElem; + private Map dependencies; + private Descriptors.Descriptor descriptor; + + public ProtobufSchemaWrapper(CustomProtobufSchema schema, MessageIndexes messageIndexes) { + this.schema = schema; + this.messageIndexes = messageIndexes; + } + + public ProtobufSchemaWrapper(CustomProtobufSchema schema, String messageTypeName) { + this.schema = schema; + this.messageTypeName = messageTypeName; + } + + + private static ProtoFileElement buildProtoFile(String schemaString) { + // Parse the schema string into a ProtoFileElement + try { + return ProtoParser.Companion.parse( + DEFAULT_LOCATION, + schemaString + ); + } catch (Exception e) { + LOGGER.log(Level.FINE, "Failed to parse schema as text format, trying binary format", e); + // Parse the binary schema into a FileDescriptorProto + try { + byte[] bytes = DECODER.decode(schemaString); + DescriptorProtos.FileDescriptorProto proto = DescriptorProtos.FileDescriptorProto.parseFrom(bytes); + return ProtobufSchemaParser.toProtoFileElement(proto); + } catch (InvalidProtocolBufferException ex) { + throw new ProtobufSchemaParser.SchemaParsingException("Failed to parse Protobuf schema in any supported format", e); + } + } + } + + /** + * Parses schema dependencies into ProtoFileElements. + * + * @param resolvedReferences Map of resolved schema dependencies + * @return Map of import paths to parsed ProtoFileElements + */ + private static Map parseDependencies(Map resolvedReferences) { + Map dependencies = new HashMap<>(); + + resolvedReferences.forEach((importPath, schemaDefinition) -> { + ProtoFileElement dependencyElement = buildProtoFile(schemaDefinition); + dependencies.put(importPath, dependencyElement); + }); + + return dependencies; + } + + /** + * Gets the descriptor for the Protobuf message. + * This method builds a descriptor from the schema definition, potentially using + * cached values for better performance. + * + * @return The Protobuf Descriptor for the message + */ + public Descriptors.Descriptor getDescriptor() { + if (descriptor == null) { + descriptor = buildDescriptor(); + } + return descriptor; + } + + private Descriptors.Descriptor buildDescriptor() { + if (schema == null || (messageIndexes == null && messageTypeName == null)) { + throw new IllegalArgumentException("Schema and message indexes must be provided"); + } + String schemaString = schema.canonicalString(); + this.rootElem = buildProtoFile(schemaString); + this.dependencies = parseDependencies(schema.resolvedReferences()); + + String messageName = toMessageName(rootElem, messageIndexes, messageTypeName); + + DynamicSchema dynamicSchema = ProtobufSchemaParser.toDynamicSchema(messageName, rootElem, dependencies); + return dynamicSchema.getMessageDescriptor(messageName); + } + + private String toMessageName(ProtoFileElement rootElem, MessageIndexes indexes, String messageTypeName) { + if (messageTypeName != null) { + return messageTypeName; + } + StringBuilder sb = new StringBuilder(); + List types = rootElem.getTypes(); + boolean first = true; + for (Integer index : indexes.getIndexes()) { + if (!first) { + sb.append("."); + } else { + first = false; + } + MessageElement message = getMessageAtIndex(types, index); + if (message == null) { + throw new IllegalArgumentException("Invalid message indexes: " + indexes); + } + sb.append(message.getName()); + types = message.getNestedTypes(); + } + return sb.toString(); + } + + private MessageElement getMessageAtIndex(List types, int index) { + int i = 0; + for (TypeElement type : types) { + if (type instanceof MessageElement) { + if (index == i) { + return (MessageElement) type; + } + i++; + } + } + return null; + } + +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/ProtobufSchemaParser.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/ProtobufSchemaParser.java new file mode 100644 index 0000000000..75f94f1c21 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/ProtobufSchemaParser.java @@ -0,0 +1,136 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse; + +import kafka.automq.table.deserializer.proto.parse.template.ProtoFileElementTemplate; +import kafka.automq.table.deserializer.proto.parse.template.ProtoSchemaFileDescriptorTemplate; +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; +import kafka.automq.table.deserializer.proto.schema.ProtobufSchema; + +import com.google.common.annotations.VisibleForTesting; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; +import com.squareup.wire.schema.internal.parser.MessageElement; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.ProtoParser; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import static kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants.DEFAULT_LOCATION; + + +public class ProtobufSchemaParser { + + /** + * Converts a ProtoFileElement and its dependencies into a DynamicSchema. + * + * @param name The name/identifier for the schema + * @param rootElem The root proto file element to convert + * @param dependencies Map of dependency file names to their ProtoFileElements + * @return A DynamicSchema representing the protobuf schema + */ + public static DynamicSchema toDynamicSchema(String name, ProtoFileElement rootElem, Map dependencies) { + Objects.requireNonNull(name, "name cannot be null"); + Objects.requireNonNull(rootElem, "rootElem cannot be null"); + Objects.requireNonNull(dependencies, "dependencies cannot be null"); + + try { + return new ProtoFileElementTemplate().convert(name, rootElem, dependencies); + } catch (Descriptors.DescriptorValidationException e) { + throw new IllegalStateException("Failed to build dynamic schema", e); + } + } + + + /** + * Converts a FileDescriptor into a ProtoFileElement. + * @param file The FileDescriptor to convert + * @return The converted ProtoFileElement + */ + public static ProtoFileElement toProtoFileElement(Descriptors.FileDescriptor file) { + return new ProtoSchemaFileDescriptorTemplate().convert(file.toProto()); + } + + + /** + * Converts a FileDescriptorProto into a ProtoFileElement. + * @param file The FileDescriptorProto to convert + * @return The converted ProtoFileElement + */ + public static ProtoFileElement toProtoFileElement(DescriptorProtos.FileDescriptorProto file) { + return new ProtoSchemaFileDescriptorTemplate().convert(file); + } + + /** + * Exception thrown when schema parsing fails. + */ + public static class SchemaParsingException extends RuntimeException { + /** + * Constructs a new schema parsing exception with the specified detail message and cause. + * + * @param message the detail message + * @param cause the cause of the exception + */ + public SchemaParsingException(String message, Throwable cause) { + super(message, cause); + } + + /** + * Constructs a new schema parsing exception with the specified detail message. + * + * @param message the detail message + */ + public SchemaParsingException(String message) { + super(message); + } + } + + @VisibleForTesting + protected static ProtobufSchema parseSchema(String schemaString, Map dependencies) { + ProtoFileElement fileElement = ProtoParser.Companion.parse( + DEFAULT_LOCATION, + schemaString + ); + Map protoDependencies = new HashMap<>(); + dependencies.forEach((importPath, schemaDefinition) -> { + ProtoFileElement dependencyElement = ProtoParser.Companion.parse( + DEFAULT_LOCATION, + schemaDefinition + ); + protoDependencies.put(importPath, dependencyElement); + }); + MessageElement firstMessage = null; + for (TypeElement typeElement : fileElement.getTypes()) { + if (typeElement instanceof MessageElement) { + firstMessage = (MessageElement) typeElement; + } + } + if (firstMessage == null) { + throw new SchemaParsingException("No message found in schema"); + } + DynamicSchema dynamicSchema = toDynamicSchema(firstMessage.getName(), fileElement, protoDependencies); + Descriptors.Descriptor messageDescriptor = dynamicSchema.getMessageDescriptor(firstMessage.getName()); + Descriptors.FileDescriptor file = messageDescriptor.getFile(); + return new ProtobufSchema(file, fileElement); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/DynamicSchemaConverter.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/DynamicSchemaConverter.java new file mode 100644 index 0000000000..1b2195e63e --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/DynamicSchemaConverter.java @@ -0,0 +1,92 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter; + +import com.squareup.wire.schema.internal.parser.OptionElement; +import com.squareup.wire.schema.internal.parser.ServiceElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.List; + +/** + * Interface defining the contract for converting source objects into components needed for DynamicSchema creation. + * This interface provides a standardized way to extract necessary information from different source types + * that can be used to build a DynamicSchema. + * + * @param The type of the source object to convert from (e.g., ProtoFileElement, FileDescriptorProto) + */ +public interface DynamicSchemaConverter { + + /** + * Extracts the syntax version from the source. + * + * @param source The source object + * @return The syntax version as a string, or null if not specified + */ + String getSyntax(T source); + + /** + * Retrieves the package name from the source. + * + * @param source The source object + * @return The package name, or null if not specified + */ + String getPackageName(T source); + + /** + * Gets all type definitions (messages and enums) from the source. + * + * @param source The source object + * @return List of TypeElements representing messages and enums + */ + List getTypes(T source); + + /** + * Retrieves the list of imports from the source. + * + * @param source The source object + * @return List of import file paths + */ + List getImports(T source); + + /** + * Retrieves the list of public imports from the source. + * + * @param source The source object + * @return List of public import file paths + */ + List getPublicImports(T source); + + /** + * Retrieves all options defined in the source. + * + * @param source The source object + * @return List of OptionElements + */ + List getOptions(T source); + + /** + * Retrieves all service definitions from the source. + * @param source The source object + * @return List of ServiceElements + */ + List getServices(T source); + +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/FileDescriptorConverter.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/FileDescriptorConverter.java new file mode 100644 index 0000000000..d27cc7468c --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/FileDescriptorConverter.java @@ -0,0 +1,431 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter; + +import kafka.automq.table.deserializer.proto.parse.converter.builder.ElementBuilder; +import kafka.automq.table.deserializer.proto.parse.converter.builder.EnumBuilder; +import kafka.automq.table.deserializer.proto.parse.converter.builder.MessageBuilder; +import kafka.automq.table.deserializer.proto.parse.converter.builder.OneOfBuilder; + +import com.google.common.collect.ImmutableList; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; +import com.squareup.wire.schema.Field; +import com.squareup.wire.schema.internal.parser.EnumElement; +import com.squareup.wire.schema.internal.parser.ExtensionsElement; +import com.squareup.wire.schema.internal.parser.FieldElement; +import com.squareup.wire.schema.internal.parser.MessageElement; +import com.squareup.wire.schema.internal.parser.OptionElement; +import com.squareup.wire.schema.internal.parser.ReservedElement; +import com.squareup.wire.schema.internal.parser.RpcElement; +import com.squareup.wire.schema.internal.parser.ServiceElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +import kotlin.ranges.IntRange; + +import static kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants.DEFAULT_LOCATION; + +/** + * Converter implementation for FileDescriptor to Wire Schema elements. + */ +public class FileDescriptorConverter implements DynamicSchemaConverter { + + @Override + public List getImports(DescriptorProtos.FileDescriptorProto source) { + List imports = new ArrayList<>(); + List dependencyList = source.getDependencyList(); + Set publicDependencyList = new HashSet<>(source.getPublicDependencyList()); + + for (int i = 0; i < dependencyList.size(); i++) { + if (!publicDependencyList.contains(i)) { + imports.add(dependencyList.get(i)); + } + } + return imports; + } + + @Override + public List getPublicImports(DescriptorProtos.FileDescriptorProto source) { + List publicImports = new ArrayList<>(); + List dependencyList = source.getDependencyList(); + List publicDependencyList = source.getPublicDependencyList(); + + for (Integer index : publicDependencyList) { + publicImports.add(dependencyList.get(index)); + } + return publicImports; + } + + @Override + public String getSyntax(DescriptorProtos.FileDescriptorProto source) { + return source.getSyntax(); + } + + @Override + public String getPackageName(DescriptorProtos.FileDescriptorProto source) { + String aPackage = source.getPackage(); + return aPackage.isEmpty() ? null : aPackage; + } + + @Override + public List getTypes(DescriptorProtos.FileDescriptorProto source) { + List types = new ArrayList<>(); + // Convert messages + for (DescriptorProtos.DescriptorProto messageType : source.getMessageTypeList()) { + types.add(convertMessage(messageType, source)); + } + + // Convert enums + for (DescriptorProtos.EnumDescriptorProto enumType : source.getEnumTypeList()) { + types.add(convertEnum(enumType)); + } + + return types; + } + + @Override + public List getServices(DescriptorProtos.FileDescriptorProto source) { + List services = new ArrayList<>(); + for (DescriptorProtos.ServiceDescriptorProto service : source.getServiceList()) { + services.add(convertService(service)); + } + return services; + } + + @Override + public List getOptions(DescriptorProtos.FileDescriptorProto source) { + List options = new ArrayList<>(); + DescriptorProtos.FileOptions fileOptions = source.getOptions(); + + // Java related options + addOptionIfPresent(options, ProtoConstants.JAVA_PACKAGE_OPTION, fileOptions.hasJavaPackage(), fileOptions.getJavaPackage()); + addOptionIfPresent(options, ProtoConstants.JAVA_OUTER_CLASSNAME_OPTION, fileOptions.hasJavaOuterClassname(), fileOptions.getJavaOuterClassname()); + addOptionIfPresent(options, ProtoConstants.JAVA_MULTIPLE_FILES_OPTION, fileOptions.hasJavaMultipleFiles(), fileOptions.getJavaMultipleFiles()); + addOptionIfPresent(options, ProtoConstants.JAVA_GENERATE_EQUALS_AND_HASH_OPTION, fileOptions.hasJavaGenerateEqualsAndHash(), fileOptions.getJavaGenerateEqualsAndHash()); + addOptionIfPresent(options, ProtoConstants.JAVA_GENERIC_SERVICES_OPTION, fileOptions.hasJavaGenericServices(), fileOptions.getJavaGenericServices()); + addOptionIfPresent(options, ProtoConstants.JAVA_STRING_CHECK_UTF8_OPTION, fileOptions.hasJavaStringCheckUtf8(), fileOptions.getJavaStringCheckUtf8()); + + // C++ related options + addOptionIfPresent(options, ProtoConstants.CC_GENERIC_SERVICES_OPTION, fileOptions.hasCcGenericServices(), fileOptions.getCcGenericServices()); + addOptionIfPresent(options, ProtoConstants.CC_ENABLE_ARENAS_OPTION, fileOptions.hasCcEnableArenas(), fileOptions.getCcEnableArenas()); + + // C# related options + addOptionIfPresent(options, ProtoConstants.CSHARP_NAMESPACE_OPTION, fileOptions.hasCsharpNamespace(), fileOptions.getCsharpNamespace()); + + // Go related options + addOptionIfPresent(options, ProtoConstants.GO_PACKAGE_OPTION, fileOptions.hasGoPackage(), fileOptions.getGoPackage()); + + // Objective-C related options + addOptionIfPresent(options, ProtoConstants.OBJC_CLASS_PREFIX_OPTION, fileOptions.hasObjcClassPrefix(), fileOptions.getObjcClassPrefix()); + + // PHP related options + addOptionIfPresent(options, ProtoConstants.PHP_CLASS_PREFIX_OPTION, fileOptions.hasPhpClassPrefix(), fileOptions.getPhpClassPrefix()); +// addOptionIfPresent(options, ProtoConstants.PHP_GENERIC_SERVICES_OPTION, fileOptions.hasPhpGenericServices(), fileOptions.getPhpGenericServices()); + addOptionIfPresent(options, ProtoConstants.PHP_METADATA_NAMESPACE_OPTION, fileOptions.hasPhpMetadataNamespace(), fileOptions.getPhpMetadataNamespace()); + addOptionIfPresent(options, ProtoConstants.PHP_NAMESPACE_OPTION, fileOptions.hasPhpNamespace(), fileOptions.getPhpNamespace()); + + // Python related options + addOptionIfPresent(options, ProtoConstants.PY_GENERIC_SERVICES_OPTION, fileOptions.hasPyGenericServices(), fileOptions.getPyGenericServices()); + + // Ruby related options + addOptionIfPresent(options, ProtoConstants.RUBY_PACKAGE_OPTION, fileOptions.hasRubyPackage(), fileOptions.getRubyPackage()); + + // Swift related options + addOptionIfPresent(options, ProtoConstants.SWIFT_PREFIX_OPTION, fileOptions.hasSwiftPrefix(), fileOptions.getSwiftPrefix()); + + // Optimize related options + addOptionIfPresent(options, ProtoConstants.OPTIMIZE_FOR_OPTION, fileOptions.hasOptimizeFor(), fileOptions.getOptimizeFor()); + + return options; + } + + private MessageElement convertMessage(DescriptorProtos.DescriptorProto descriptor, + DescriptorProtos.FileDescriptorProto file) { + MessageBuilder builder = new MessageBuilder(descriptor.getName()); + + // Add options + new MessageOptionStrategy().addOption(builder, descriptor.getOptions()); + + // Process fields and oneofs + Map oneofBuilders = new HashMap<>(); + for (DescriptorProtos.OneofDescriptorProto oneof : descriptor.getOneofDeclList()) { + oneofBuilders.put(oneofBuilders.size(), builder.newOneOf(oneof.getName())); + } + + // Process fields + for (DescriptorProtos.FieldDescriptorProto field : descriptor.getFieldList()) { + FieldElement fieldElement = convertField(file, field, field.hasOneofIndex()); + if (field.hasOneofIndex()) { + oneofBuilders.get(field.getOneofIndex()).addField(fieldElement); + } else { + builder.addField(fieldElement); + } + } + + // Process nested types + for (DescriptorProtos.DescriptorProto nestedType : descriptor.getNestedTypeList()) { + builder.addNestedType(convertMessage(nestedType, file)); + } + for (DescriptorProtos.EnumDescriptorProto enumType : descriptor.getEnumTypeList()) { + builder.addNestedType(convertEnum(enumType)); + } + + // Process reserved ranges and names + processReservedRanges(descriptor, builder); + processExtensionRanges(descriptor, builder); + + return builder.build(); + } + + private void processReservedRanges(DescriptorProtos.DescriptorProto descriptor, MessageBuilder builder) { + for (String reservedName : descriptor.getReservedNameList()) { + builder.addReserved(new ReservedElement(DEFAULT_LOCATION, "", Collections.singletonList(reservedName))); + } + + for (DescriptorProtos.DescriptorProto.ReservedRange range : descriptor.getReservedRangeList()) { + List values = new ArrayList<>(); + values.add(new IntRange(range.getStart(), range.getEnd() - 1)); + builder.addReserved(new ReservedElement(DEFAULT_LOCATION, "", values)); + } + } + + private void processExtensionRanges(DescriptorProtos.DescriptorProto descriptor, MessageBuilder builder) { + for (DescriptorProtos.DescriptorProto.ExtensionRange range : descriptor.getExtensionRangeList()) { + List values = new ArrayList<>(); + values.add(new IntRange(range.getStart(), range.getEnd() - 1)); + builder.addExtension(new ExtensionsElement(DEFAULT_LOCATION, "", values)); + } + } + + private EnumElement convertEnum(DescriptorProtos.EnumDescriptorProto enumType) { + EnumBuilder builder = new EnumBuilder(enumType.getName()); + + // Add options + new EnumOptionStrategy().addOption(builder, enumType.getOptions()); + + // Add constants + for (DescriptorProtos.EnumValueDescriptorProto value : enumType.getValueList()) { + builder.addConstant(value.getName(), value.getNumber()); + } + + // Process reserved ranges and names + for (DescriptorProtos.EnumDescriptorProto.EnumReservedRange range : enumType.getReservedRangeList()) { + List values = new ArrayList<>(); + values.add(new IntRange(range.getStart(), range.getEnd() - 1)); + builder.addReserved(new ReservedElement(DEFAULT_LOCATION, "", values)); + } + + for (String reservedName : enumType.getReservedNameList()) { + builder.addReserved(new ReservedElement(DEFAULT_LOCATION, "", Collections.singletonList(reservedName))); + } + + return builder.build(); + } + + private ServiceElement convertService(DescriptorProtos.ServiceDescriptorProto sv) { + String name = sv.getName(); + ImmutableList.Builder rpcs = ImmutableList.builder(); + + // Convert methods + for (DescriptorProtos.MethodDescriptorProto md : sv.getMethodList()) { + // Build method options + ImmutableList.Builder methodOptions = ImmutableList.builder(); + + if (md.getOptions().hasDeprecated()) { + methodOptions.add(new OptionElement(ProtoConstants.DEPRECATED_OPTION, OptionElement.Kind.BOOLEAN, + md.getOptions().getDeprecated(), false)); + } + + if (md.getOptions().hasIdempotencyLevel()) { + methodOptions.add(new OptionElement(ProtoConstants.IDEMPOTENCY_LEVEL_OPTION, OptionElement.Kind.ENUM, + md.getOptions().getIdempotencyLevel(), false)); + } + + // Create RPC element + rpcs.add(new RpcElement(DEFAULT_LOCATION, md.getName(), "", + getTypeName(md.getInputType()), + getTypeName(md.getOutputType()), + md.getClientStreaming(), + md.getServerStreaming(), + methodOptions.build())); + } + + // Build service options + ImmutableList.Builder serviceOptions = ImmutableList.builder(); + if (sv.getOptions().hasDeprecated()) { + serviceOptions.add(new OptionElement(ProtoConstants.DEPRECATED_OPTION, OptionElement.Kind.BOOLEAN, + sv.getOptions().getDeprecated(), false)); + } + + return new ServiceElement(DEFAULT_LOCATION, name, "", rpcs.build(), serviceOptions.build()); + } + + /** + * Determines the field label based on the proto syntax version and field properties. + */ + private Field.Label label(DescriptorProtos.FileDescriptorProto file, DescriptorProtos.FieldDescriptorProto fd) { + boolean isProto3 = file.getSyntax().equals(ProtoConstants.PROTO3); + switch (fd.getLabel()) { + case LABEL_REQUIRED: + return isProto3 ? null : Field.Label.REQUIRED; + case LABEL_OPTIONAL: + // If it's a Proto3 optional, we have to print the optional label. + return isProto3 && !fd.hasProto3Optional() ? null : Field.Label.OPTIONAL; + case LABEL_REPEATED: + return Field.Label.REPEATED; + default: + throw new IllegalArgumentException("Unsupported label"); + } + } + + /** + * Returns the field type name, either from the type name or the primitive type. + */ + private String dataType(DescriptorProtos.FieldDescriptorProto field) { + if (field.hasTypeName()) { + return field.getTypeName(); + } else { + DescriptorProtos.FieldDescriptorProto.Type type = field.getType(); + return Descriptors.FieldDescriptor.Type.valueOf(type).name().toLowerCase(Locale.ENGLISH); + } + } + + private FieldElement convertField(DescriptorProtos.FileDescriptorProto file, + DescriptorProtos.FieldDescriptorProto fd, boolean inOneof) { + String name = fd.getName(); + DescriptorProtos.FieldOptions fieldDescriptorOptions = fd.getOptions(); + List optionElements = new ArrayList<>(); + // Add standard field options if present + if (fd.hasJsonName() && !fd.getJsonName().equals(getDefaultJsonName(name))) { + optionElements.add(new OptionElement(ProtoConstants.JSON_NAME_OPTION, OptionElement.Kind.STRING, + fd.getJsonName(), false)); + } + addOptionIfPresent(optionElements, ProtoConstants.PACKED_OPTION, fieldDescriptorOptions.hasPacked(), fd.getOptions().getPacked()); + addOptionIfPresent(optionElements, ProtoConstants.DEPRECATED_OPTION, fieldDescriptorOptions.hasDeprecated(), fieldDescriptorOptions.getDeprecated()); + addOptionIfPresent(optionElements, ProtoConstants.CTYPE_OPTION, fieldDescriptorOptions.hasCtype(), fieldDescriptorOptions.getCtype()); + addOptionIfPresent(optionElements, ProtoConstants.JSTYPE_OPTION, fieldDescriptorOptions.hasJstype(), fieldDescriptorOptions.getJstype()); + ImmutableList.Builder options = ImmutableList.builder(); + options.addAll(optionElements); + + String jsonName = null; // Let Wire calculate the default JSON name + String defaultValue = fd.hasDefaultValue() && fd.getDefaultValue() != null ? fd.getDefaultValue() : null; + + return new FieldElement( + DEFAULT_LOCATION, + inOneof ? null : label(file, fd), + dataType(fd), + name, + defaultValue, + jsonName, + fd.getNumber(), + "", + options.build() + ); + } + + /** + * Strategy interface for handling different types of options. + */ + private interface OptionStrategy { + void addOption(ElementBuilder builder, Object options); + } + + /** + * Strategy for handling message options. + */ + private static class MessageOptionStrategy implements OptionStrategy { + @Override + public void addOption(ElementBuilder builder, Object options) { + DescriptorProtos.MessageOptions msgOptions = (DescriptorProtos.MessageOptions) options; + if (msgOptions.hasMapEntry()) { + builder.addOption(ProtoConstants.MAP_ENTRY_OPTION, OptionElement.Kind.BOOLEAN, msgOptions.getMapEntry()); + } + if (msgOptions.hasNoStandardDescriptorAccessor()) { + builder.addOption(ProtoConstants.NO_STANDARD_DESCRIPTOR_OPTION, OptionElement.Kind.BOOLEAN, + msgOptions.getNoStandardDescriptorAccessor()); + } + } + } + + /** + * Strategy for handling enum options. + */ + private static class EnumOptionStrategy implements OptionStrategy { + @Override + public void addOption(ElementBuilder builder, Object options) { + DescriptorProtos.EnumOptions enumOptions = (DescriptorProtos.EnumOptions) options; + if (enumOptions.hasAllowAlias()) { + builder.addOption(ProtoConstants.ALLOW_ALIAS_OPTION, OptionElement.Kind.BOOLEAN, enumOptions.getAllowAlias()); + } + } + } + + /** + * Adds an option to the options list if it is present in the source. + * Determines the appropriate OptionElement.Kind based on the value type. + * + * @param options The list of options to add to + * @param name The name of the option + * @param hasOption Whether the option is present + * @param value The value of the option + */ + private void addOptionIfPresent(List options, String name, boolean hasOption, Object value) { + if (hasOption) { + OptionElement.Kind kind; + if (value instanceof Boolean) { + kind = OptionElement.Kind.BOOLEAN; + } else if (value instanceof String) { + kind = OptionElement.Kind.STRING; + } else if (value instanceof Enum) { + kind = OptionElement.Kind.ENUM; + } else { + kind = OptionElement.Kind.STRING; + } + options.add(new OptionElement(name, kind, value.toString(), false)); + } + } + + private String getTypeName(String typeName) { + return typeName.startsWith(".") ? typeName : "." + typeName; + } + + /** + * Calculates the default JSON name for a field following the protobuf convention. + * Converts from snake_case to lowerCamelCase. + */ + private String getDefaultJsonName(String fieldName) { + String[] parts = fieldName.split("_"); + StringBuilder defaultJsonName = new StringBuilder(parts[0]); + for (int i = 1; i < parts.length; ++i) { + defaultJsonName.append(parts[i].substring(0, 1).toUpperCase(Locale.ENGLISH)) + .append(parts[i].substring(1)); + } + return defaultJsonName.toString(); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoConstants.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoConstants.java new file mode 100644 index 0000000000..14a9ae18d7 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoConstants.java @@ -0,0 +1,169 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter; + +import com.google.protobuf.AnyProto; +import com.google.protobuf.ApiProto; +import com.google.protobuf.Descriptors; +import com.google.protobuf.DurationProto; +import com.google.protobuf.EmptyProto; +import com.google.protobuf.FieldMaskProto; +import com.google.protobuf.SourceContextProto; +import com.google.protobuf.StructProto; +import com.google.protobuf.TimestampProto; +import com.google.protobuf.TypeProto; +import com.google.protobuf.WrappersProto; +import com.google.type.CalendarPeriodProto; +import com.google.type.ColorProto; +import com.google.type.DateProto; +import com.google.type.DayOfWeek; +import com.google.type.ExprProto; +import com.google.type.FractionProto; +import com.google.type.IntervalProto; +import com.google.type.LatLng; +import com.google.type.LocalizedTextProto; +import com.google.type.MoneyProto; +import com.google.type.MonthProto; +import com.google.type.PhoneNumberProto; +import com.google.type.PostalAddressProto; +import com.google.type.QuaternionProto; +import com.google.type.TimeOfDayProto; +import com.squareup.wire.schema.Location; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; + +import java.util.Arrays; +import java.util.Map; +import java.util.stream.Collectors; + +import static kafka.automq.table.deserializer.proto.parse.ProtobufSchemaParser.toProtoFileElement; + +/** + * Constants used in protobuf schema parsing and conversion. + */ +public final class ProtoConstants { + /** Default location for proto file elements */ + public static final Location DEFAULT_LOCATION = Location.get(""); + + /** Protocol buffer syntax versions */ + public static final String PROTO2 = "proto2"; + public static final String PROTO3 = "proto3"; + + /** Common protobuf option names */ + public static final String ALLOW_ALIAS_OPTION = "allow_alias"; + public static final String MAP_ENTRY_OPTION = "map_entry"; + public static final String DEPRECATED_OPTION = "deprecated"; + public static final String PACKED_OPTION = "packed"; + public static final String JSON_NAME_OPTION = "json_name"; + public static final String CTYPE_OPTION = "ctype"; + public static final String JSTYPE_OPTION = "jstype"; + + /** Map field related constants */ + public static final String KEY_FIELD = "key"; + public static final String VALUE_FIELD = "value"; + public static final String MAP_ENTRY_SUFFIX = "Entry"; + + /** Message options */ + public static final String NO_STANDARD_DESCRIPTOR_OPTION = "no_standard_descriptor_accessor"; + + /** RPC options */ + public static final String IDEMPOTENCY_LEVEL_OPTION = "idempotency_level"; + + /** Java related options */ + public static final String JAVA_PACKAGE_OPTION = "java_package"; + public static final String JAVA_OUTER_CLASSNAME_OPTION = "java_outer_classname"; + public static final String JAVA_MULTIPLE_FILES_OPTION = "java_multiple_files"; + public static final String JAVA_GENERATE_EQUALS_AND_HASH_OPTION = "java_generate_equals_and_hash"; + public static final String JAVA_GENERIC_SERVICES_OPTION = "java_generic_services"; + public static final String JAVA_STRING_CHECK_UTF8_OPTION = "java_string_check_utf8"; + + /** C++ related options */ + public static final String CC_GENERIC_SERVICES_OPTION = "cc_generic_services"; + public static final String CC_ENABLE_ARENAS_OPTION = "cc_enable_arenas"; + + /** C# related options */ + public static final String CSHARP_NAMESPACE_OPTION = "csharp_namespace"; + + /** Go related options */ + public static final String GO_PACKAGE_OPTION = "go_package"; + + /** Objective-C related options */ + public static final String OBJC_CLASS_PREFIX_OPTION = "objc_class_prefix"; + + /** PHP related options */ + public static final String PHP_CLASS_PREFIX_OPTION = "php_class_prefix"; + public static final String PHP_GENERIC_SERVICES_OPTION = "php_generic_services"; + public static final String PHP_METADATA_NAMESPACE_OPTION = "php_metadata_namespace"; + public static final String PHP_NAMESPACE_OPTION = "php_namespace"; + + /** Python related options */ + public static final String PY_GENERIC_SERVICES_OPTION = "py_generic_services"; + + /** Ruby related options */ + public static final String RUBY_PACKAGE_OPTION = "ruby_package"; + + /** Swift related options */ + public static final String SWIFT_PREFIX_OPTION = "swift_prefix"; + + /** Optimize related options */ + public static final String OPTIMIZE_FOR_OPTION = "optimize_for"; + + + /** Well-known protobuf type descriptors */ + public static final Descriptors.FileDescriptor[] WELL_KNOWN_DEPENDENCIES; + + public static final Map BASE_DEPENDENCIES; + + static { + // Support all the Protobuf WellKnownTypes + // and the protos from Google API, https://github.com/googleapis/googleapis + WELL_KNOWN_DEPENDENCIES = new Descriptors.FileDescriptor[] { + ApiProto.getDescriptor().getFile(), + FieldMaskProto.getDescriptor().getFile(), + SourceContextProto.getDescriptor().getFile(), + StructProto.getDescriptor().getFile(), + TypeProto.getDescriptor().getFile(), + TimestampProto.getDescriptor().getFile(), + WrappersProto.getDescriptor().getFile(), + AnyProto.getDescriptor().getFile(), + EmptyProto.getDescriptor().getFile(), + DurationProto.getDescriptor().getFile(), + TimeOfDayProto.getDescriptor().getFile(), + DateProto.getDescriptor().getFile(), + CalendarPeriodProto.getDescriptor().getFile(), + ColorProto.getDescriptor().getFile(), + DayOfWeek.getDescriptor().getFile(), + LatLng.getDescriptor().getFile(), + FractionProto.getDescriptor().getFile(), + MoneyProto.getDescriptor().getFile(), + MonthProto.getDescriptor().getFile(), + PhoneNumberProto.getDescriptor().getFile(), + PostalAddressProto.getDescriptor().getFile(), + LocalizedTextProto.getDescriptor().getFile(), + IntervalProto.getDescriptor().getFile(), + ExprProto.getDescriptor().getFile(), + QuaternionProto.getDescriptor().getFile(), + }; + + BASE_DEPENDENCIES = Arrays.stream(WELL_KNOWN_DEPENDENCIES) + .collect(Collectors.toMap( + Descriptors.FileDescriptor::getFullName, + item -> toProtoFileElement(item.toProto()))); + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoElementConvert.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoElementConvert.java new file mode 100644 index 0000000000..2b0c895b2e --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoElementConvert.java @@ -0,0 +1,35 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter; + +import kafka.automq.table.deserializer.proto.schema.EnumDefinition; +import kafka.automq.table.deserializer.proto.schema.MessageDefinition; + +import com.squareup.wire.schema.internal.parser.EnumElement; +import com.squareup.wire.schema.internal.parser.MessageElement; + +/** + * Visitor interface for processing protobuf elements. + * Implementations can provide different ways to process these elements. + */ +public interface ProtoElementConvert { + EnumDefinition convert(EnumElement enumElement); + MessageDefinition convert(MessageElement messageElement); +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoElementSchemaConvert.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoElementSchemaConvert.java new file mode 100644 index 0000000000..5af84686d9 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoElementSchemaConvert.java @@ -0,0 +1,187 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter; + +import kafka.automq.table.deserializer.proto.schema.EnumDefinition; +import kafka.automq.table.deserializer.proto.schema.MessageDefinition; + +import com.squareup.wire.schema.Field; +import com.squareup.wire.schema.ProtoType; +import com.squareup.wire.schema.internal.parser.EnumConstantElement; +import com.squareup.wire.schema.internal.parser.EnumElement; +import com.squareup.wire.schema.internal.parser.FieldElement; +import com.squareup.wire.schema.internal.parser.MessageElement; +import com.squareup.wire.schema.internal.parser.OneOfElement; +import com.squareup.wire.schema.internal.parser.OptionElement; +import com.squareup.wire.schema.internal.parser.ReservedElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Optional; +import java.util.Set; + +import kotlin.ranges.IntRange; + +public class ProtoElementSchemaConvert implements ProtoElementConvert { + + @Override + public EnumDefinition convert(EnumElement enumElement) { + Boolean allowAlias = getOptionBoolean(ProtoConstants.ALLOW_ALIAS_OPTION, enumElement.getOptions()); + EnumDefinition.Builder builder = EnumDefinition.newBuilder(enumElement.getName(), allowAlias); + + for (EnumConstantElement constant : enumElement.getConstants()) { + builder.addValue(constant.getName(), constant.getTag()); + } + + return builder.build(); + } + + @Override + public MessageDefinition convert(MessageElement messageElement) { + MessageDefinition.Builder message = MessageDefinition.newBuilder(messageElement.getName()); + + // Process nested types + for (TypeElement type : messageElement.getNestedTypes()) { + if (type instanceof MessageElement) { + message.addMessageDefinition(convert((MessageElement) type)); + } else if (type instanceof EnumElement) { + message.addEnumDefinition(convert((EnumElement) type)); + } + } + + // Process fields + processMessageFields(message, messageElement); + + // Process reserved ranges and names + processReservedElements(message, messageElement); + + // Process options + processMessageOptions(message, messageElement); + + return message.build(); + } + + private void processMessageFields(MessageDefinition.Builder message, MessageElement messageElement) { + Set processedFields = new HashSet<>(); + + // Process oneofs first + for (OneOfElement oneof : messageElement.getOneOfs()) { + MessageDefinition.OneofBuilder oneofBuilder = message.addOneof(oneof.getName()); + for (FieldElement field : oneof.getFields()) { + processedFields.add(field.getName()); + addFieldToOneof(oneofBuilder, field); + } + } + + // Process regular fields + for (FieldElement field : messageElement.getFields()) { + if (!processedFields.contains(field.getName())) { + addFieldToMessage(message, field); + } + } + } + + private void addFieldToMessage(MessageDefinition.Builder message, FieldElement field) { + Field.Label fieldLabel = field.getLabel(); + String label = fieldLabel != null ? fieldLabel.toString().toLowerCase(Locale.ENGLISH) : null; + String fieldType = field.getType(); + ProtoType protoType = ProtoType.get(fieldType); + + // Handle map fields + if (protoType.isMap()) { + ProtoType keyType = protoType.getKeyType(); + ProtoType valueType = protoType.getValueType(); + if (keyType != null && valueType != null) { + processMapField(message, field, keyType, valueType); + return; + } + } + + message.addField(label, fieldType, field.getName(), field.getTag(), + field.getDefaultValue(), field.getJsonName(), + getOptionBoolean(ProtoConstants.PACKED_OPTION, field.getOptions())); + } + + private void processMapField(MessageDefinition.Builder message, FieldElement field, + ProtoType keyType, ProtoType valueType) { + String mapEntryName = toMapEntryName(field.getName()); + MessageDefinition.Builder mapMessage = MessageDefinition.newBuilder(mapEntryName); + mapMessage.setMapEntry(true); + + mapMessage.addField(null, resolveFieldTypeName(keyType), ProtoConstants.KEY_FIELD, 1, null, null, null); + mapMessage.addField(null, resolveFieldTypeName(valueType), ProtoConstants.VALUE_FIELD, 2, null, null, null); + + message.addMessageDefinition(mapMessage.build()); + message.addField("repeated", mapEntryName, field.getName(), field.getTag(), + null, field.getJsonName(), null); + } + + private void addFieldToOneof(MessageDefinition.OneofBuilder oneof, FieldElement field) { + oneof.addField(field.getType(), field.getName(), field.getTag(), + field.getDefaultValue(), field.getJsonName()); + } + + private void processReservedElements(MessageDefinition.Builder message, MessageElement messageElement) { + for (ReservedElement reserved : messageElement.getReserveds()) { + for (Object value : reserved.getValues()) { + if (value instanceof String) { + message.addReservedName((String) value); + } else if (value instanceof Integer) { + int tag = (Integer) value; + message.addReservedRange(tag, tag); + } else if (value instanceof IntRange) { + IntRange range = (IntRange) value; + message.addReservedRange(range.getStart(), range.getEndInclusive()); + } + } + } + } + + private void processMessageOptions(MessageDefinition.Builder message, MessageElement messageElement) { + Boolean isMapEntry = getOptionBoolean(ProtoConstants.MAP_ENTRY_OPTION, messageElement.getOptions()); + if (isMapEntry != null) { + message.setMapEntry(isMapEntry); + } + } + + private static Boolean getOptionBoolean(String name, List options) { + return findOption(name, options) + .map(o -> Boolean.valueOf(o.getValue().toString())) + .orElse(null); + } + + private static Optional findOption(String name, List options) { + return options.stream() + .filter(o -> o.getName().equals(name)) + .findFirst(); + } + + private static String toMapEntryName(String fieldName) { + return Character.toUpperCase(fieldName.charAt(0)) + + fieldName.substring(1) + + ProtoConstants.MAP_ENTRY_SUFFIX; + } + + private static String resolveFieldTypeName(ProtoType type) { + return type.toString(); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoFileElementConverter.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoFileElementConverter.java new file mode 100644 index 0000000000..ae6a3854f4 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/ProtoFileElementConverter.java @@ -0,0 +1,73 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter; + +import com.squareup.wire.Syntax; +import com.squareup.wire.schema.internal.parser.OptionElement; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.ServiceElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.List; + +/** + * Implementation of DynamicSchemaConverter for ProtoFileElement source objects. + * This class provides specific conversion logic for extracting schema information + * from ProtoFileElement instances. + */ +public class ProtoFileElementConverter implements DynamicSchemaConverter { + + @Override + public String getSyntax(ProtoFileElement source) { + Syntax syntax = source.getSyntax(); + return syntax != null ? syntax.toString() : null; + } + + @Override + public String getPackageName(ProtoFileElement source) { + String packageName = source.getPackageName(); + return packageName != null ? packageName : ""; + } + + @Override + public List getTypes(ProtoFileElement source) { + return source.getTypes(); + } + + @Override + public List getImports(ProtoFileElement source) { + return source.getImports(); + } + + @Override + public List getPublicImports(ProtoFileElement source) { + return source.getPublicImports(); + } + + @Override + public List getOptions(ProtoFileElement source) { + return source.getOptions(); + } + + @Override + public List getServices(ProtoFileElement source) { + return source.getServices(); + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/ElementBuilder.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/ElementBuilder.java new file mode 100644 index 0000000000..a0c6d1b9eb --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/ElementBuilder.java @@ -0,0 +1,50 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter.builder; + +import com.google.common.collect.ImmutableList; +import com.squareup.wire.schema.Location; +import com.squareup.wire.schema.internal.parser.OptionElement; + +/** + * Abstract builder for converting protobuf elements to Wire Schema elements. + */ +public abstract class ElementBuilder> { + protected static final Location DEFAULT_LOCATION = Location.get(""); + protected static final String DOCUMENTATION = ""; + protected final String name; + protected final ImmutableList.Builder options = ImmutableList.builder(); + + protected ElementBuilder(String name) { + this.name = name; + } + + protected abstract T build(); + + @SuppressWarnings("unchecked") + protected B self() { + return (B) this; + } + + public B addOption(String name, OptionElement.Kind kind, Object value) { + options.add(new OptionElement(name, kind, value.toString(), false)); + return self(); + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/EnumBuilder.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/EnumBuilder.java new file mode 100644 index 0000000000..3638312e7f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/EnumBuilder.java @@ -0,0 +1,65 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter.builder; + +import com.google.common.collect.ImmutableList; +import com.squareup.wire.schema.internal.parser.EnumConstantElement; +import com.squareup.wire.schema.internal.parser.EnumElement; +import com.squareup.wire.schema.internal.parser.ReservedElement; + +/** + * Builder for EnumElement construction. + */ +public class EnumBuilder extends ElementBuilder { + private final ImmutableList.Builder constants = ImmutableList.builder(); + private final ImmutableList.Builder reserved = ImmutableList.builder(); + + public EnumBuilder(String name) { + super(name); + } + + public EnumBuilder addConstant(String name, int number) { + constants.add(new EnumConstantElement( + DEFAULT_LOCATION, + name, + number, + DOCUMENTATION, + ImmutableList.of() + )); + return this; + } + + public EnumBuilder addReserved(ReservedElement element) { + reserved.add(element); + return this; + } + + @Override + public EnumElement build() { + return new EnumElement( + DEFAULT_LOCATION, + name, + DOCUMENTATION, + options.build(), + constants.build(), + reserved.build() + ); + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/MessageBuilder.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/MessageBuilder.java new file mode 100644 index 0000000000..273aa2d0ec --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/MessageBuilder.java @@ -0,0 +1,93 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter.builder; + +import com.google.common.collect.ImmutableList; +import com.squareup.wire.schema.internal.parser.ExtensionsElement; +import com.squareup.wire.schema.internal.parser.FieldElement; +import com.squareup.wire.schema.internal.parser.MessageElement; +import com.squareup.wire.schema.internal.parser.ReservedElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +/** + * Builder for MessageElement construction. + */ +public class MessageBuilder extends ElementBuilder { + private final ImmutableList.Builder fields = ImmutableList.builder(); + private final ImmutableList.Builder nested = ImmutableList.builder(); + private final ImmutableList.Builder reserved = ImmutableList.builder(); + private final ImmutableList.Builder extensions = ImmutableList.builder(); + private final List oneofs = new ArrayList<>(); + + public MessageBuilder(String name) { + super(name); + } + + public MessageBuilder addField(FieldElement field) { + fields.add(field); + return this; + } + + public MessageBuilder addNestedType(TypeElement type) { + nested.add(type); + return this; + } + + public MessageBuilder addReserved(ReservedElement element) { + reserved.add(element); + return this; + } + + public MessageBuilder addExtension(ExtensionsElement element) { + extensions.add(element); + return this; + } + + public OneOfBuilder newOneOf(String name) { + OneOfBuilder builder = new OneOfBuilder(name); + oneofs.add(builder); + return builder; + } + + @Override + public MessageElement build() { + return new MessageElement( + DEFAULT_LOCATION, + name, + DOCUMENTATION, + nested.build(), + options.build(), + reserved.build(), + fields.build(), + oneofs.stream() + .filter(b -> !b.getFields().isEmpty()) + .map(OneOfBuilder::build) + .collect(Collectors.toList()), + extensions.build(), + Collections.emptyList(), + Collections.emptyList() + ); + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/OneOfBuilder.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/OneOfBuilder.java new file mode 100644 index 0000000000..0ee084779c --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/converter/builder/OneOfBuilder.java @@ -0,0 +1,58 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.converter.builder; + +import com.google.common.collect.ImmutableList; +import com.squareup.wire.schema.internal.parser.FieldElement; +import com.squareup.wire.schema.internal.parser.OneOfElement; + +import java.util.Collections; + +/** + * Builder for OneOfElement construction. + */ +public class OneOfBuilder extends ElementBuilder { + private final ImmutableList.Builder fields = ImmutableList.builder(); + + public OneOfBuilder(String name) { + super(name); + } + + public OneOfBuilder addField(FieldElement field) { + fields.add(field); + return this; + } + + public ImmutableList getFields() { + return fields.build(); + } + + @Override + protected OneOfElement build() { + return new OneOfElement( + name, + DOCUMENTATION, + fields.build(), + Collections.emptyList(), + Collections.emptyList(), + DEFAULT_LOCATION + ); + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/DynamicSchemaTemplate.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/DynamicSchemaTemplate.java new file mode 100644 index 0000000000..b15a4c4a2e --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/DynamicSchemaTemplate.java @@ -0,0 +1,173 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.template; + +import kafka.automq.table.deserializer.proto.parse.converter.DynamicSchemaConverter; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoElementSchemaConvert; +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; + +import com.google.protobuf.Descriptors; +import com.squareup.wire.schema.internal.parser.EnumElement; +import com.squareup.wire.schema.internal.parser.MessageElement; +import com.squareup.wire.schema.internal.parser.OptionElement; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants.JAVA_MULTIPLE_FILES_OPTION; +import static kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants.JAVA_OUTER_CLASSNAME_OPTION; +import static kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants.JAVA_PACKAGE_OPTION; + + +/** + * Abstract template class that defines the skeleton algorithm for converting source objects into DynamicSchema. + * This class implements the Template Method pattern, providing a structured approach to schema conversion + * while allowing specific steps to be customized by subclasses. + * + * @param The type of the source object to convert from + */ +public abstract class DynamicSchemaTemplate { + + /** + * Gets the converter implementation for the specific source type. + * + * @return The converter implementation + */ + protected abstract DynamicSchemaConverter getConverter(); + + + /** + * Processes schema options and applies them to the schema builder. + * + * @param schema The schema builder + * @param options The list of options to process + */ + protected void processSchemaOptions(DynamicSchema.Builder schema, List options) { + findOption(JAVA_PACKAGE_OPTION, options) + .ifPresent(o -> schema.setJavaPackage(o.getValue().toString())); + + findOption(JAVA_OUTER_CLASSNAME_OPTION, options) + .ifPresent(o -> schema.setJavaOuterClassname(o.getValue().toString())); + + findOption(JAVA_MULTIPLE_FILES_OPTION, options) + .ifPresent(o -> schema.setJavaMultipleFiles(Boolean.valueOf(o.getValue().toString()))); + } + + /** + * Finds an option by name in a list of options. + * + * @param name The option name + * @param options The list of options + * @return Optional containing the found option, or empty if not found + */ + public static Optional findOption(String name, List options) { + return options.stream() + .filter(o -> o.getName().equals(name)) + .findFirst(); + } + + /** + * Converts the source object into a DynamicSchema. + * This method implements the template method pattern, defining the skeleton of the conversion algorithm + * while delegating specific conversion steps to the converter. + * + * @param name The name of the schema + * @param source The source object to convert + * @param dependencies Map of dependencies + * @return The converted DynamicSchema + * @throws Descriptors.DescriptorValidationException if validation fails + */ + public DynamicSchema convert(String name, T source, Map dependencies) + throws Descriptors.DescriptorValidationException { + + DynamicSchemaConverter converter = getConverter(); + ProtoElementSchemaConvert elementSchemaConvert = new ProtoElementSchemaConvert(); + DynamicSchema.Builder schema = DynamicSchema.newBuilder(); + + // Set basic properties + Optional.ofNullable(converter.getSyntax(source)) + .ifPresent(schema::setSyntax); + + Optional.ofNullable(converter.getPackageName(source)) + .ifPresent(schema::setPackage); + + // Process types + for (TypeElement typeElem : converter.getTypes(source)) { + if (typeElem instanceof MessageElement) { + MessageElement messageElement = (MessageElement) typeElem; + schema.addMessageDefinition(elementSchemaConvert.convert(messageElement)); + } else if (typeElem instanceof EnumElement) { + EnumElement enumElement = (EnumElement) typeElem; + schema.addEnumDefinition(elementSchemaConvert.convert(enumElement)); + } + } + + // Process imports and dependencies + processImportsAndDependencies(schema, source, converter, dependencies); + + // Process options + processSchemaOptions(schema, converter.getOptions(source)); + + schema.setName(name); + return schema.build(); + } + + /** + * Processes imports and dependencies, adding them to the schema builder. + * + * @param schema The schema builder + * @param source The source object + * @param converter The converter instance + * @param dependencies Map of dependencies + */ + protected void processImportsAndDependencies(DynamicSchema.Builder schema, T source, + DynamicSchemaConverter converter, + Map dependencies) { + // Process regular imports + for (String ref : converter.getImports(source)) { + ProtoFileElement dep = dependencies.get(ref) == null ? ProtoConstants.BASE_DEPENDENCIES.get(ref) : dependencies.get(ref); + if (dep != null) { + schema.addDependency(ref); + try { + schema.addSchema(convert(ref, (T) dep, dependencies)); + } catch (Descriptors.DescriptorValidationException e) { + throw new IllegalStateException("Failed to process import: " + ref, e); + } + } + } + + // Process public imports + for (String ref : converter.getPublicImports(source)) { + ProtoFileElement dep = dependencies.get(ref) == null ? ProtoConstants.BASE_DEPENDENCIES.get(ref) : dependencies.get(ref); + if (dep != null) { + schema.addPublicDependency(ref); + try { + schema.addSchema(convert(ref, (T) dep, dependencies)); + } catch (Descriptors.DescriptorValidationException e) { + throw new IllegalStateException("Failed to process public import: " + ref, e); + } + } + } + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoFileElementTemplate.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoFileElementTemplate.java new file mode 100644 index 0000000000..d29a0616c9 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoFileElementTemplate.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.template; + +import kafka.automq.table.deserializer.proto.parse.converter.DynamicSchemaConverter; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoFileElementConverter; + +import com.squareup.wire.schema.internal.parser.ProtoFileElement; + +/** + * Template implementation for converting ProtoFileElement to DynamicSchema. + * This class provides the concrete implementation of the template methods + * specific to ProtoFileElement processing. + */ +public class ProtoFileElementTemplate extends DynamicSchemaTemplate { + + private final ProtoFileElementConverter converter; + + /** + * Creates a new ProtoFileElementTemplate with a default converter. + */ + public ProtoFileElementTemplate() { + this.converter = new ProtoFileElementConverter(); + } + + @Override + protected DynamicSchemaConverter getConverter() { + return converter; + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoSchemaFileDescriptorTemplate.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoSchemaFileDescriptorTemplate.java new file mode 100644 index 0000000000..ce0fd65187 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoSchemaFileDescriptorTemplate.java @@ -0,0 +1,32 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.template; + +import kafka.automq.table.deserializer.proto.parse.converter.DynamicSchemaConverter; +import kafka.automq.table.deserializer.proto.parse.converter.FileDescriptorConverter; + +import com.google.protobuf.DescriptorProtos; + +public class ProtoSchemaFileDescriptorTemplate extends ProtoSchemaTemplate { + @Override + protected DynamicSchemaConverter getConverter() { + return new FileDescriptorConverter(); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoSchemaTemplate.java b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoSchemaTemplate.java new file mode 100644 index 0000000000..d020b35710 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/parse/template/ProtoSchemaTemplate.java @@ -0,0 +1,76 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse.template; + +import kafka.automq.table.deserializer.proto.parse.converter.DynamicSchemaConverter; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants; + +import com.google.common.collect.ImmutableList; +import com.squareup.wire.Syntax; +import com.squareup.wire.schema.internal.parser.OptionElement; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.ServiceElement; +import com.squareup.wire.schema.internal.parser.TypeElement; + +import java.util.Collections; + +/** + * Template method for converting protobuf schema representations + */ +public abstract class ProtoSchemaTemplate { + + protected abstract DynamicSchemaConverter getConverter(); + + public ProtoFileElement convert(T source) { + DynamicSchemaConverter converter = getConverter(); + + Syntax syntax = ProtoConstants.PROTO3.equals(converter.getSyntax(source)) ? + Syntax.PROTO_3 : Syntax.PROTO_2; + + String packageName = converter.getPackageName(source); + + ImmutableList.Builder imports = ImmutableList.builder(); + imports.addAll(converter.getImports(source)); + + ImmutableList.Builder publicImports = ImmutableList.builder(); + publicImports.addAll(converter.getPublicImports(source)); + + ImmutableList.Builder types = ImmutableList.builder(); + types.addAll(converter.getTypes(source)); + + ImmutableList.Builder services = ImmutableList.builder(); + services.addAll(converter.getServices(source)); + + ImmutableList.Builder options = ImmutableList.builder(); + options.addAll(converter.getOptions(source)); + + return new ProtoFileElement( + ProtoConstants.DEFAULT_LOCATION, + packageName, + syntax, + imports.build(), + publicImports.build(), + types.build(), + services.build(), + Collections.emptyList(), // extends + options.build() + ); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/schema/DynamicSchema.java b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/DynamicSchema.java new file mode 100644 index 0000000000..10b324d7c4 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/DynamicSchema.java @@ -0,0 +1,447 @@ +/* + * Copyright 2021 Red Hat + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.deserializer.proto.schema; + +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; + +public class DynamicSchema { + // --- public static --- + + /** + * Creates a new dynamic schema builder + * + * @return the schema builder + */ + public static Builder newBuilder() { + return new Builder(); + } + + /** + * Parses a serialized schema descriptor (from input stream; closes the stream) + * + * @param schemaDescIn the descriptor input stream + * @return the schema object + */ + public static DynamicSchema parseFrom(InputStream schemaDescIn) + throws Descriptors.DescriptorValidationException, IOException { + try { + int len; + byte[] buf = new byte[4096]; + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + while ((len = schemaDescIn.read(buf)) > 0) { + baos.write(buf, 0, len); + } + return parseFrom(baos.toByteArray()); + } finally { + schemaDescIn.close(); + } + } + + /** + * Parses a serialized schema descriptor (from byte array) + * + * @param schemaDescBuf the descriptor byte array + * @return the schema object + */ + public static DynamicSchema parseFrom(byte[] schemaDescBuf) + throws Descriptors.DescriptorValidationException, IOException { + return new DynamicSchema(DescriptorProtos.FileDescriptorSet.parseFrom(schemaDescBuf)); + } + + // --- public --- + + /** + * Gets the protobuf file descriptor proto + * + * @return the file descriptor proto + */ + public DescriptorProtos.FileDescriptorProto getFileDescriptorProto() { + return mFileDescSet.getFile(0); + } + + /** + * Creates a new dynamic message builder for the given message type + * + * @param msgTypeName the message type name + * @return the message builder (null if not found) + */ + public DynamicMessage.Builder newMessageBuilder(String msgTypeName) { + Descriptors.Descriptor msgType = getMessageDescriptor(msgTypeName); + if (msgType == null) { + return null; + } + return DynamicMessage.newBuilder(msgType); + } + + /** + * Gets the protobuf message descriptor for the given message type + * + * @param msgTypeName the message type name + * @return the message descriptor (null if not found) + */ + public Descriptors.Descriptor getMessageDescriptor(String msgTypeName) { + Descriptors.Descriptor msgType = mMsgDescriptorMapShort.get(msgTypeName); + if (msgType == null) { + msgType = mMsgDescriptorMapFull.get(msgTypeName); + } + return msgType; + } + + /** + * Gets the enum value for the given enum type and name + * + * @param enumTypeName the enum type name + * @param enumName the enum name + * @return the enum value descriptor (null if not found) + */ + public Descriptors.EnumValueDescriptor getEnumValue(String enumTypeName, String enumName) { + Descriptors.EnumDescriptor enumType = getEnumDescriptor(enumTypeName); + if (enumType == null) { + return null; + } + return enumType.findValueByName(enumName); + } + + /** + * Gets the enum value for the given enum type and number + * + * @param enumTypeName the enum type name + * @param enumNumber the enum number + * @return the enum value descriptor (null if not found) + */ + public Descriptors.EnumValueDescriptor getEnumValue(String enumTypeName, int enumNumber) { + Descriptors.EnumDescriptor enumType = getEnumDescriptor(enumTypeName); + if (enumType == null) { + return null; + } + return enumType.findValueByNumber(enumNumber); + } + + /** + * Gets the protobuf enum descriptor for the given enum type + * + * @param enumTypeName the enum type name + * @return the enum descriptor (null if not found) + */ + public Descriptors.EnumDescriptor getEnumDescriptor(String enumTypeName) { + Descriptors.EnumDescriptor enumType = mEnumDescriptorMapShort.get(enumTypeName); + if (enumType == null) { + enumType = mEnumDescriptorMapFull.get(enumTypeName); + } + return enumType; + } + + /** + * Returns the message types registered with the schema + * + * @return the set of message type names + */ + public Set getMessageTypes() { + return new TreeSet(mMsgDescriptorMapFull.keySet()); + } + + /** + * Returns the enum types registered with the schema + * + * @return the set of enum type names + */ + public Set getEnumTypes() { + return new TreeSet(mEnumDescriptorMapFull.keySet()); + } + + /** + * Serializes the schema + * + * @return the serialized schema descriptor + */ + public byte[] toByteArray() { + return mFileDescSet.toByteArray(); + } + + /** + * Returns a string representation of the schema + * + * @return the schema string + */ + public String toString() { + Set msgTypes = getMessageTypes(); + Set enumTypes = getEnumTypes(); + return "types: " + msgTypes + "\nenums: " + enumTypes + "\n" + mFileDescSet; + } + + // --- private --- + + private DynamicSchema(DescriptorProtos.FileDescriptorSet fileDescSet) + throws Descriptors.DescriptorValidationException { + mFileDescSet = fileDescSet; + Map fileDescMap = init(fileDescSet); + + Set msgDupes = new HashSet(); + Set enumDupes = new HashSet(); + for (Descriptors.FileDescriptor fileDesc : fileDescMap.values()) { + for (Descriptors.Descriptor msgType : fileDesc.getMessageTypes()) { + addMessageType(msgType, null, msgDupes, enumDupes); + } + for (Descriptors.EnumDescriptor enumType : fileDesc.getEnumTypes()) { + addEnumType(enumType, null, enumDupes); + } + } + + for (String msgName : msgDupes) { + mMsgDescriptorMapShort.remove(msgName); + } + for (String enumName : enumDupes) { + mEnumDescriptorMapShort.remove(enumName); + } + } + + @SuppressWarnings("unchecked") + private Map init(DescriptorProtos.FileDescriptorSet fileDescSet) + throws Descriptors.DescriptorValidationException { + // check for dupes + Set allFdProtoNames = new HashSet(); + for (DescriptorProtos.FileDescriptorProto fdProto : fileDescSet.getFileList()) { + if (allFdProtoNames.contains(fdProto.getName())) { + throw new IllegalArgumentException("duplicate name: " + fdProto.getName()); + } + allFdProtoNames.add(fdProto.getName()); + } + + // build FileDescriptors, resolve dependencies (imports) if any + Map resolvedFileDescMap = new HashMap(); + while (resolvedFileDescMap.size() < fileDescSet.getFileCount()) { + for (DescriptorProtos.FileDescriptorProto fdProto : fileDescSet.getFileList()) { + if (resolvedFileDescMap.containsKey(fdProto.getName())) { + continue; + } + + // getDependencyList() signature was changed and broke compatibility in 2.6.1; workaround + // with reflection + // List dependencyList = fdProto.getDependencyList(); + List dependencyList = null; + try { + Method m = fdProto.getClass().getMethod("getDependencyList", (Class[]) null); + dependencyList = (List) m.invoke(fdProto, (Object[]) null); + } catch (Exception e) { + throw new RuntimeException(e); + } + + List resolvedFdList = new ArrayList(); + for (String depName : dependencyList) { + if (!allFdProtoNames.contains(depName)) { + throw new IllegalArgumentException( + "cannot resolve import " + depName + " in " + fdProto.getName()); + } + Descriptors.FileDescriptor fd = resolvedFileDescMap.get(depName); + if (fd != null) { + resolvedFdList.add(fd); + } + } + + if (resolvedFdList.size() == dependencyList.size()) { // dependencies resolved + Descriptors.FileDescriptor[] fds = new Descriptors.FileDescriptor[resolvedFdList.size()]; + Descriptors.FileDescriptor fd = Descriptors.FileDescriptor.buildFrom(fdProto, + resolvedFdList.toArray(fds)); + resolvedFileDescMap.put(fdProto.getName(), fd); + } + } + } + + return resolvedFileDescMap; + } + + private void addMessageType(Descriptors.Descriptor msgType, String scope, Set msgDupes, + Set enumDupes) { + String msgTypeNameFull = msgType.getFullName(); + String msgTypeNameShort = scope == null ? msgType.getName() : scope + "." + msgType.getName(); + + if (mMsgDescriptorMapFull.containsKey(msgTypeNameFull)) { + throw new IllegalArgumentException("duplicate name: " + msgTypeNameFull); + } + if (mMsgDescriptorMapShort.containsKey(msgTypeNameShort)) { + msgDupes.add(msgTypeNameShort); + } + + mMsgDescriptorMapFull.put(msgTypeNameFull, msgType); + mMsgDescriptorMapShort.put(msgTypeNameShort, msgType); + + for (Descriptors.Descriptor nestedType : msgType.getNestedTypes()) { + addMessageType(nestedType, msgTypeNameShort, msgDupes, enumDupes); + } + for (Descriptors.EnumDescriptor enumType : msgType.getEnumTypes()) { + addEnumType(enumType, msgTypeNameShort, enumDupes); + } + } + + private void addEnumType(Descriptors.EnumDescriptor enumType, String scope, Set enumDupes) { + String enumTypeNameFull = enumType.getFullName(); + String enumTypeNameShort = scope == null ? enumType.getName() : scope + "." + enumType.getName(); + + if (mEnumDescriptorMapFull.containsKey(enumTypeNameFull)) { + throw new IllegalArgumentException("duplicate name: " + enumTypeNameFull); + } + if (mEnumDescriptorMapShort.containsKey(enumTypeNameShort)) { + enumDupes.add(enumTypeNameShort); + } + + mEnumDescriptorMapFull.put(enumTypeNameFull, enumType); + mEnumDescriptorMapShort.put(enumTypeNameShort, enumType); + } + + private DescriptorProtos.FileDescriptorSet mFileDescSet; + private Map mMsgDescriptorMapFull = new HashMap(); + private Map mMsgDescriptorMapShort = new HashMap(); + private Map mEnumDescriptorMapFull = new HashMap(); + private Map mEnumDescriptorMapShort = new HashMap(); + + /** + * DynamicSchema.Builder + */ + public static class Builder { + // --- public --- + + /** + * Builds a dynamic schema + * + * @return the schema object + */ + public DynamicSchema build() throws Descriptors.DescriptorValidationException { + DescriptorProtos.FileDescriptorSet.Builder fileDescSetBuilder = DescriptorProtos.FileDescriptorSet + .newBuilder(); + fileDescSetBuilder.addFile(mFileDescProtoBuilder.build()); + fileDescSetBuilder.mergeFrom(mFileDescSetBuilder.build()); + return new DynamicSchema(fileDescSetBuilder.build()); + } + + public Builder setSyntax(String syntax) { + mFileDescProtoBuilder.setSyntax(syntax); + return this; + } + + public Builder setName(String name) { + // if name does not end with ".proto", append it + if (!name.endsWith(".proto")) { + name += ".proto"; + } + + mFileDescProtoBuilder.setName(name); + return this; + } + + public Builder setPackage(String name) { + mFileDescProtoBuilder.setPackage(name); + return this; + } + + public Builder addMessageDefinition(MessageDefinition msgDef) { + mFileDescProtoBuilder.addMessageType(msgDef.getMessageType()); + return this; + } + + public Builder addEnumDefinition(EnumDefinition enumDef) { + mFileDescProtoBuilder.addEnumType(enumDef.getEnumType()); + return this; + } + + // Note: added + public Builder addDependency(String dependency) { + mFileDescProtoBuilder.addDependency(dependency); + return this; + } + + // Note: added + public Builder addPublicDependency(String dependency) { + for (int i = 0; i < mFileDescProtoBuilder.getDependencyCount(); i++) { + if (mFileDescProtoBuilder.getDependency(i).equals(dependency)) { + mFileDescProtoBuilder.addPublicDependency(i); + return this; + } + } + mFileDescProtoBuilder.addDependency(dependency); + mFileDescProtoBuilder.addPublicDependency(mFileDescProtoBuilder.getDependencyCount() - 1); + return this; + } + + // Note: added + public Builder setJavaPackage(String javaPackage) { + DescriptorProtos.FileOptions.Builder optionsBuilder = DescriptorProtos.FileOptions.newBuilder(); + optionsBuilder.setJavaPackage(javaPackage); + mFileDescProtoBuilder.mergeOptions(optionsBuilder.build()); + return this; + } + + // Note: added + public Builder setJavaOuterClassname(String javaOuterClassname) { + DescriptorProtos.FileOptions.Builder optionsBuilder = DescriptorProtos.FileOptions.newBuilder(); + optionsBuilder.setJavaOuterClassname(javaOuterClassname); + mFileDescProtoBuilder.mergeOptions(optionsBuilder.build()); + return this; + } + + // Note: added + public Builder setJavaMultipleFiles(boolean javaMultipleFiles) { + DescriptorProtos.FileOptions.Builder optionsBuilder = DescriptorProtos.FileOptions.newBuilder(); + optionsBuilder.setJavaMultipleFiles(javaMultipleFiles); + mFileDescProtoBuilder.mergeOptions(optionsBuilder.build()); + return this; + } + + // Note: changed + public Builder addSchema(DynamicSchema schema) { + for (DescriptorProtos.FileDescriptorProto file : schema.mFileDescSet.getFileList()) { + if (!contains(file)) { + mFileDescSetBuilder.addFile(file); + } + } + return this; + } + + // Note: added + private boolean contains(DescriptorProtos.FileDescriptorProto fileDesc) { + List files = mFileDescSetBuilder.getFileList(); + for (DescriptorProtos.FileDescriptorProto file : files) { + if (file.getName().equals(fileDesc.getName())) { + return true; + } + } + return false; + } + + // --- private --- + + private Builder() { + mFileDescProtoBuilder = DescriptorProtos.FileDescriptorProto.newBuilder(); + mFileDescSetBuilder = DescriptorProtos.FileDescriptorSet.newBuilder(); + } + + private DescriptorProtos.FileDescriptorProto.Builder mFileDescProtoBuilder; + private DescriptorProtos.FileDescriptorSet.Builder mFileDescSetBuilder; + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/schema/EnumDefinition.java b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/EnumDefinition.java new file mode 100644 index 0000000000..36c97c16a4 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/EnumDefinition.java @@ -0,0 +1,84 @@ +/* + * Copyright 2021 Red Hat + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.deserializer.proto.schema; + +import com.google.protobuf.DescriptorProtos; + +public class EnumDefinition { + // --- public static --- + + public static Builder newBuilder(String enumName) { + return newBuilder(enumName, null); + } + + public static Builder newBuilder(String enumName, Boolean allowAlias) { + return new Builder(enumName, allowAlias); + } + + // --- public --- + + public String toString() { + return mEnumType.toString(); + } + + // --- package --- + + DescriptorProtos.EnumDescriptorProto getEnumType() { + return mEnumType; + } + + // --- private --- + + private EnumDefinition(DescriptorProtos.EnumDescriptorProto enumType) { + mEnumType = enumType; + } + + private DescriptorProtos.EnumDescriptorProto mEnumType; + + /** + * EnumDefinition.Builder + */ + public static class Builder { + // --- public --- + + public Builder addValue(String name, int num) { + DescriptorProtos.EnumValueDescriptorProto.Builder enumValBuilder = DescriptorProtos.EnumValueDescriptorProto + .newBuilder(); + enumValBuilder.setName(name).setNumber(num); + mEnumTypeBuilder.addValue(enumValBuilder.build()); + return this; + } + + public EnumDefinition build() { + return new EnumDefinition(mEnumTypeBuilder.build()); + } + + // --- private --- + + private Builder(String enumName, Boolean allowAlias) { + mEnumTypeBuilder = DescriptorProtos.EnumDescriptorProto.newBuilder(); + mEnumTypeBuilder.setName(enumName); + if (allowAlias != null) { + DescriptorProtos.EnumOptions.Builder optionsBuilder = DescriptorProtos.EnumOptions + .newBuilder(); + optionsBuilder.setAllowAlias(allowAlias); + mEnumTypeBuilder.mergeOptions(optionsBuilder.build()); + } + } + + private DescriptorProtos.EnumDescriptorProto.Builder mEnumTypeBuilder; + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/schema/MessageDefinition.java b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/MessageDefinition.java new file mode 100644 index 0000000000..afe220d43b --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/MessageDefinition.java @@ -0,0 +1,221 @@ +/* + * Copyright 2021 Red Hat + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.deserializer.proto.schema; + +import com.google.protobuf.DescriptorProtos; + +import java.util.HashMap; +import java.util.Map; + +public class MessageDefinition { + // --- public static --- + + public static Builder newBuilder(String msgTypeName) { + return new Builder(msgTypeName); + } + + // --- public --- + + public String toString() { + return mMsgType.toString(); + } + + // --- package --- + + DescriptorProtos.DescriptorProto getMessageType() { + return mMsgType; + } + + // --- private --- + + private MessageDefinition(DescriptorProtos.DescriptorProto msgType) { + mMsgType = msgType; + } + + private DescriptorProtos.DescriptorProto mMsgType; + + /** + * MessageDefinition.Builder + */ + public static class Builder { + // --- public --- + + public Builder addField(String label, String type, String name, int num, String defaultVal) { + return addField(label, type, name, num, defaultVal, null, null); + } + + public Builder addField(String label, String type, String name, int num, String defaultVal, + String jsonName, Boolean isPacked) { + DescriptorProtos.FieldDescriptorProto.Label protoLabel = sLabelMap.get(label); + doAddField(protoLabel, type, name, num, defaultVal, jsonName, isPacked, null); + return this; + } + + public OneofBuilder addOneof(String oneofName) { + mMsgTypeBuilder.addOneofDecl( + DescriptorProtos.OneofDescriptorProto.newBuilder().setName(oneofName).build()); + return new OneofBuilder(this, mOneofIndex++); + } + + public Builder addMessageDefinition(MessageDefinition msgDef) { + mMsgTypeBuilder.addNestedType(msgDef.getMessageType()); + return this; + } + + public Builder addEnumDefinition(EnumDefinition enumDef) { + mMsgTypeBuilder.addEnumType(enumDef.getEnumType()); + return this; + } + + // Note: added + public Builder addReservedName(String reservedName) { + mMsgTypeBuilder.addReservedName(reservedName); + return this; + } + + // Note: added + public Builder addReservedRange(int start, int end) { + DescriptorProtos.DescriptorProto.ReservedRange.Builder rangeBuilder = DescriptorProtos.DescriptorProto.ReservedRange + .newBuilder(); + rangeBuilder.setStart(start).setEnd(end); + mMsgTypeBuilder.addReservedRange(rangeBuilder.build()); + return this; + } + + // Note: added + public Builder setMapEntry(boolean mapEntry) { + DescriptorProtos.MessageOptions.Builder optionsBuilder = DescriptorProtos.MessageOptions + .newBuilder(); + optionsBuilder.setMapEntry(mapEntry); + mMsgTypeBuilder.mergeOptions(optionsBuilder.build()); + return this; + } + + public MessageDefinition build() { + return new MessageDefinition(mMsgTypeBuilder.build()); + } + + // --- private --- + + private Builder(String msgTypeName) { + mMsgTypeBuilder = DescriptorProtos.DescriptorProto.newBuilder(); + mMsgTypeBuilder.setName(msgTypeName); + } + + private void doAddField(DescriptorProtos.FieldDescriptorProto.Label label, String type, String name, + int num, String defaultVal, String jsonName, Boolean isPacked, OneofBuilder oneofBuilder) { + DescriptorProtos.FieldDescriptorProto.Builder fieldBuilder = DescriptorProtos.FieldDescriptorProto + .newBuilder(); + // Note: changed + if (label != null) { + fieldBuilder.setLabel(label); + } + DescriptorProtos.FieldDescriptorProto.Type primType = sTypeMap.get(type); + if (primType != null) { + fieldBuilder.setType(primType); + } else { + fieldBuilder.setTypeName(type); + } + fieldBuilder.setName(name).setNumber(num); + if (defaultVal != null) { + fieldBuilder.setDefaultValue(defaultVal); + } + if (oneofBuilder != null) { + fieldBuilder.setOneofIndex(oneofBuilder.getIdx()); + } + if (jsonName != null) { + fieldBuilder.setJsonName(jsonName); + } + if (isPacked != null) { + DescriptorProtos.FieldOptions.Builder optionsBuilder = DescriptorProtos.FieldOptions + .newBuilder(); + optionsBuilder.setPacked(isPacked); + fieldBuilder.mergeOptions(optionsBuilder.build()); + } + mMsgTypeBuilder.addField(fieldBuilder.build()); + } + + private DescriptorProtos.DescriptorProto.Builder mMsgTypeBuilder; + private int mOneofIndex = 0; + } + + /** + * MessageDefinition.OneofBuilder + */ + public static class OneofBuilder { + // --- public --- + + public OneofBuilder addField(String type, String name, int num, String defaultVal) { + return addField(type, name, num, defaultVal, null); + } + + public OneofBuilder addField(String type, String name, int num, String defaultVal, String jsonName) { + mMsgBuilder.doAddField(DescriptorProtos.FieldDescriptorProto.Label.LABEL_OPTIONAL, type, name, + num, defaultVal, jsonName, null, this); + return this; + } + + public Builder msgDefBuilder() { + return mMsgBuilder; + } + + public int getIdx() { + return mIdx; + } + + // --- private --- + + private OneofBuilder(Builder msgBuilder, int oneofIdx) { + mMsgBuilder = msgBuilder; + mIdx = oneofIdx; + } + + private Builder mMsgBuilder; + private int mIdx; + } + + // --- private static --- + + private static Map sTypeMap; + private static Map sLabelMap; + + static { + sTypeMap = new HashMap(); + sTypeMap.put("double", DescriptorProtos.FieldDescriptorProto.Type.TYPE_DOUBLE); + sTypeMap.put("float", DescriptorProtos.FieldDescriptorProto.Type.TYPE_FLOAT); + sTypeMap.put("int32", DescriptorProtos.FieldDescriptorProto.Type.TYPE_INT32); + sTypeMap.put("int64", DescriptorProtos.FieldDescriptorProto.Type.TYPE_INT64); + sTypeMap.put("uint32", DescriptorProtos.FieldDescriptorProto.Type.TYPE_UINT32); + sTypeMap.put("uint64", DescriptorProtos.FieldDescriptorProto.Type.TYPE_UINT64); + sTypeMap.put("sint32", DescriptorProtos.FieldDescriptorProto.Type.TYPE_SINT32); + sTypeMap.put("sint64", DescriptorProtos.FieldDescriptorProto.Type.TYPE_SINT64); + sTypeMap.put("fixed32", DescriptorProtos.FieldDescriptorProto.Type.TYPE_FIXED32); + sTypeMap.put("fixed64", DescriptorProtos.FieldDescriptorProto.Type.TYPE_FIXED64); + sTypeMap.put("sfixed32", DescriptorProtos.FieldDescriptorProto.Type.TYPE_SFIXED32); + sTypeMap.put("sfixed64", DescriptorProtos.FieldDescriptorProto.Type.TYPE_SFIXED64); + sTypeMap.put("bool", DescriptorProtos.FieldDescriptorProto.Type.TYPE_BOOL); + sTypeMap.put("string", DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING); + sTypeMap.put("bytes", DescriptorProtos.FieldDescriptorProto.Type.TYPE_BYTES); + // sTypeMap.put("enum", FieldDescriptorProto.Type.TYPE_ENUM); + // sTypeMap.put("message", FieldDescriptorProto.Type.TYPE_MESSAGE); + // sTypeMap.put("group", FieldDescriptorProto.Type.TYPE_GROUP); + + sLabelMap = new HashMap(); + sLabelMap.put("optional", DescriptorProtos.FieldDescriptorProto.Label.LABEL_OPTIONAL); + sLabelMap.put("required", DescriptorProtos.FieldDescriptorProto.Label.LABEL_REQUIRED); + sLabelMap.put("repeated", DescriptorProtos.FieldDescriptorProto.Label.LABEL_REPEATED); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/schema/MessageIndexes.java b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/MessageIndexes.java new file mode 100644 index 0000000000..a8bdc9a6b1 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/MessageIndexes.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.schema; + +import org.apache.kafka.common.utils.ByteUtils; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +public class MessageIndexes { + private static final List DEFAULT_INDEX = Collections.singletonList(0); + private final List indexes; + + public MessageIndexes(List indexes) { + this.indexes = new ArrayList<>(indexes); + } + + public List getIndexes() { + return Collections.unmodifiableList(indexes); + } + + public static MessageIndexes readFrom(ByteBuffer buffer) { + int size = ByteUtils.readVarint(buffer); + if (size == 0) { + return new MessageIndexes(DEFAULT_INDEX); + } + List indexes = new ArrayList<>(size); + for (int i = 0; i < size; i++) { + indexes.add(ByteUtils.readVarint(buffer)); + } + return new MessageIndexes(indexes); + } +} diff --git a/core/src/main/java/kafka/automq/table/deserializer/proto/schema/ProtobufSchema.java b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/ProtobufSchema.java new file mode 100644 index 0000000000..107fc56a59 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/deserializer/proto/schema/ProtobufSchema.java @@ -0,0 +1,49 @@ +/* + * Copyright 2021 Red Hat + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.deserializer.proto.schema; + +import com.google.protobuf.Descriptors.FileDescriptor; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; + +import java.util.Objects; + +public class ProtobufSchema { + + private final FileDescriptor fileDescriptor; + private ProtoFileElement protoFileElement; + + public ProtobufSchema(FileDescriptor fileDescriptor, ProtoFileElement protoFileElement) { + Objects.requireNonNull(fileDescriptor); + Objects.requireNonNull(protoFileElement); + this.fileDescriptor = fileDescriptor; + this.protoFileElement = protoFileElement; + } + + /** + * @return the fileDescriptor + */ + public FileDescriptor getFileDescriptor() { + return fileDescriptor; + } + + /** + * @return the protoFileElement + */ + public ProtoFileElement getProtoFileElement() { + return protoFileElement; + } + +} diff --git a/core/src/main/java/kafka/automq/table/events/AvroCodec.java b/core/src/main/java/kafka/automq/table/events/AvroCodec.java new file mode 100644 index 0000000000..2ba4cd7317 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/AvroCodec.java @@ -0,0 +1,85 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.Schema; +import org.apache.avro.generic.IndexedRecord; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.BinaryEncoder; +import org.apache.avro.io.DatumReader; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.EncoderFactory; +import org.apache.avro.specific.SpecificDatumReader; +import org.apache.avro.specific.SpecificDatumWriter; +import org.apache.iceberg.avro.CodecSetup; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; + +public class AvroCodec { + private static final byte[] MAGIC_BYTES = new byte[] {(byte) 0x23, (byte) 0x33}; + + static { + CodecSetup.setup(); + } + + public static byte[] encode(T data) throws IOException { + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + DataOutputStream dataOut = new DataOutputStream(out); + + // Write the magic bytes + dataOut.write(MAGIC_BYTES); + + // Write avro schema + dataOut.writeUTF(data.getSchema().toString()); + BinaryEncoder encoder = EncoderFactory.get().binaryEncoder(out, null); + DatumWriter writer = new SpecificDatumWriter<>(data.getSchema()); + writer.write(data, encoder); + encoder.flush(); + return out.toByteArray(); + } + } + + public static T decode(byte[] data) throws IOException { + try (ByteArrayInputStream in = new ByteArrayInputStream(data, 0, data.length)) { + DataInputStream dataInput = new DataInputStream(in); + + // Read the magic bytes + byte header0 = dataInput.readByte(); + byte header1 = dataInput.readByte(); + if (header0 != MAGIC_BYTES[0] || header1 != MAGIC_BYTES[1]) { + throw new IllegalArgumentException(String.format("Invalid magic bytes: 0x%02X%02X", header0, header1)); + } + + // Read avro schema + Schema avroSchema = new Schema.Parser().parse(dataInput.readUTF()); + + // Decode the datum with the parsed avro schema. + BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(in, null); + DatumReader reader = new SpecificDatumReader<>(avroSchema); + reader.setSchema(avroSchema); + return reader.read(null, decoder); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/events/CommitRequest.java b/core/src/main/java/kafka/automq/table/events/CommitRequest.java new file mode 100644 index 0000000000..62dcb89b7d --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/CommitRequest.java @@ -0,0 +1,130 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.util.Utf8; + +import java.util.List; +import java.util.UUID; + +public class CommitRequest implements Payload { + public static final int NOOP_SPEC_ID = -1; + private UUID commitId; + private String topic; + private List offsets; + private int specId; + private final Schema avroSchema; + + private static final Schema AVRO_SCHEMA = SchemaBuilder.builder().record(CommitRequest.class.getName()) + .fields() + .name("commitId").type(UUID_SCHEMA).noDefault() + .name("topic").type().stringType().noDefault() + .name("offsets").type().array().items(WorkerOffset.AVRO_SCHEMA).noDefault() + .name("specId").type().nullable().intType().intDefault(NOOP_SPEC_ID) + .endRecord(); + + // used by avro deserialize reflection + public CommitRequest(Schema schema) { + this.avroSchema = schema; + } + + public CommitRequest(UUID commitId, String topic, List offsets) { + this(commitId, topic, NOOP_SPEC_ID, offsets); + } + + public CommitRequest(UUID commitId, String topic, int specId, List offsets) { + this.commitId = commitId; + this.topic = topic; + this.offsets = offsets; + this.specId = specId; + this.avroSchema = AVRO_SCHEMA; + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.commitId = Element.toUuid((GenericData.Fixed) v); + return; + case 1: + this.topic = ((Utf8) v).toString(); + return; + case 2: + //noinspection unchecked + this.offsets = (List) v; + return; + case 3: + this.specId = (Integer) v; + return; + default: + // ignore the object, it must be from a newer version of the format + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return Element.toFixed(commitId); + case 1: + return topic; + case 2: + return offsets; + case 3: + return specId; + default: + throw new IllegalArgumentException("Unknown field index: " + i); + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + public UUID commitId() { + return commitId; + } + + public String topic() { + return topic; + } + + public List offsets() { + return offsets; + } + + public int specId() { + return specId; + } + + @Override + public String toString() { + return "CommitRequest{" + + "commitId=" + commitId + + ", topic='" + topic + '\'' + + ", offsets=" + offsets + + ", specId=" + specId + + '}'; + } +} diff --git a/core/src/main/java/kafka/automq/table/events/CommitResponse.java b/core/src/main/java/kafka/automq/table/events/CommitResponse.java new file mode 100644 index 0000000000..7458124d7e --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/CommitResponse.java @@ -0,0 +1,287 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.kafka.common.utils.ByteBufferInputStream; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.BinaryEncoder; +import org.apache.avro.io.DatumReader; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.EncoderFactory; +import org.apache.avro.util.Utf8; +import org.apache.commons.io.input.BoundedInputStream; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.avro.GenericAvroReader; +import org.apache.iceberg.avro.GenericAvroWriter; +import org.apache.iceberg.types.Types; + +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +public class CommitResponse implements Payload { + private static final byte ICEBERG_MAGIC = 0x01; + + private int code; + private UUID commitId; + private String topic; + private List nextOffsets; + private List dataFiles; + private List deleteFiles; + private TopicMetric topicMetric = TopicMetric.NOOP; + private List partitionMetrics = Collections.emptyList(); + private final Schema avroSchema; + private Schema dataFileSchema; + private Schema deleteFileSchema; + + // used by avro deserialize reflection + public CommitResponse(Schema schema) { + this.avroSchema = schema; + } + + public CommitResponse(Types.StructType partitionType, int code, UUID commitId, String topic, + List nextOffsets, List dataFiles, List deleteFiles, TopicMetric topicMetric, List partitionMetrics) { + this.code = code; + this.commitId = commitId; + this.topic = topic; + this.nextOffsets = nextOffsets; + this.dataFiles = dataFiles; + this.deleteFiles = deleteFiles; + this.topicMetric = topicMetric; + this.partitionMetrics = partitionMetrics; + + Types.StructType dataFileStruct = DataFile.getType(partitionType); + + Map dataFileNames = new HashMap<>(); + dataFileNames.put(dataFileStruct, "org.apache.iceberg.GenericDataFile"); + dataFileNames.put(partitionType, "org.apache.iceberg.PartitionData"); + this.dataFileSchema = AvroSchemaUtil.convert(dataFileStruct, dataFileNames); + + Map deleteFileNames = new HashMap<>(); + deleteFileNames.put(dataFileStruct, "org.apache.iceberg.GenericDeleteFile"); + deleteFileNames.put(partitionType, "org.apache.iceberg.PartitionData"); + this.deleteFileSchema = AvroSchemaUtil.convert(dataFileStruct, deleteFileNames); + + this.avroSchema = SchemaBuilder.builder().record(CommitResponse.class.getName()) + .fields() + .name("code").type().intType().noDefault() + .name("commitId").type(UUID_SCHEMA).noDefault() + .name("topic").type().stringType().noDefault() + .name("nextOffsets").type().array().items(WorkerOffset.AVRO_SCHEMA).noDefault() + .name("dataFiles").type().bytesType().noDefault() + .name("deleteFiles").type().bytesType().noDefault() + .name("topicMetric").type(TopicMetric.AVRO_SCHEMA).withDefault(TopicMetric.NOOP) + .name("partitionMetrics").type().array().items(PartitionMetric.AVRO_SCHEMA).noDefault() + .endRecord(); + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.code = (int) v; + break; + case 1: + this.commitId = Element.toUuid((GenericData.Fixed) v); + break; + case 2: + this.topic = ((Utf8) v).toString(); + break; + case 3: + //noinspection unchecked + this.nextOffsets = (List) v; + break; + case 4: + try { + this.dataFiles = decodeIcebergArray((ByteBuffer) v); + } catch (IOException e) { + throw new RuntimeException(e); + } + break; + case 5: + try { + this.deleteFiles = decodeIcebergArray((ByteBuffer) v); + } catch (IOException e) { + throw new RuntimeException(e); + } + break; + case 6: + this.topicMetric = (TopicMetric) v; + break; + case 7: + //noinspection unchecked + this.partitionMetrics = (List) v; + break; + default: + // ignore the object, it must be from a newer version of the format + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return code; + case 1: + return Element.toFixed(commitId); + case 2: + return topic; + case 3: + return nextOffsets; + case 4: + try { + return ByteBuffer.wrap(encodeIcebergArray(dataFiles, dataFileSchema)); + } catch (IOException e) { + throw new RuntimeException(e); + } + case 5: + try { + return ByteBuffer.wrap(encodeIcebergArray(deleteFiles, deleteFileSchema)); + } catch (IOException e) { + throw new RuntimeException(e); + } + case 6: + return topicMetric; + case 7: + return partitionMetrics; + default: + throw new IllegalArgumentException("Unknown field index: " + i); + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + public int code() { + return code; + } + + public UUID commitId() { + return commitId; + } + + public String topic() { + return topic; + } + + public List nextOffsets() { + return nextOffsets; + } + + public List dataFiles() { + return dataFiles; + } + + public List deleteFiles() { + return deleteFiles; + } + + public TopicMetric topicMetric() { + return topicMetric; + } + + public List partitionMetrics() { + return partitionMetrics; + } + + @Override + public String toString() { + return "CommitResponse{" + + "code=" + code + + ", commitId=" + commitId + + ", topic='" + topic + '\'' + + ", nextOffsets=" + nextOffsets + + ", dataFiles=" + dataFiles.size() + + ", deleteFiles=" + deleteFiles.size() + + ", topicMetric=" + topicMetric + + ", partitionMetrics=" + partitionMetrics + + '}'; + } + + byte[] encodeIcebergArray(List list, Schema schema) throws IOException { + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + DataOutputStream dataOut = new DataOutputStream(out); + + // Write the magic bytes + dataOut.write(ICEBERG_MAGIC); + + // Write avro schema + dataOut.writeUTF(schema.toString()); + + // Encode the datum with avro schema. + DatumWriter writer = GenericAvroWriter.create(schema); + for (T datum : list) { + ByteArrayOutputStream elementOut = new ByteArrayOutputStream(); + BinaryEncoder encoder = EncoderFactory.get().binaryEncoder(elementOut, null); + writer.write(datum, encoder); + encoder.flush(); + byte[] elementOutBytes = elementOut.toByteArray(); + dataOut.writeInt(elementOutBytes.length); + dataOut.write(elementOutBytes); + } + + return out.toByteArray(); + } + } + + List decodeIcebergArray(ByteBuffer data) throws IOException { + try ( + ByteBufferInputStream in = new ByteBufferInputStream(data); + DataInputStream dataInput = new DataInputStream(in) + ) { + // Read the magic bytes + byte magic = dataInput.readByte(); + if (magic != ICEBERG_MAGIC) { + throw new IllegalArgumentException(String.format("Unrecognized magic byte: 0x%02X", magic)); + } + + // Read avro schema + Schema avroSchema = new Schema.Parser().parse(dataInput.readUTF()); + + List list = new ArrayList<>(); + // Decode the datum with the parsed avro schema. + DatumReader reader = GenericAvroReader.create(avroSchema); + reader.setSchema(avroSchema); + while (in.available() != 0) { + int length = dataInput.readInt(); + BinaryDecoder binaryDecoder = DecoderFactory.get().binaryDecoder(BoundedInputStream.builder().setInputStream(in).setMaxCount(length).get(), null); + list.add(reader.read(null, binaryDecoder)); + } + return list; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/events/Element.java b/core/src/main/java/kafka/automq/table/events/Element.java new file mode 100644 index 0000000000..949bab278d --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/Element.java @@ -0,0 +1,49 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.IndexedRecord; +import org.apache.avro.specific.SpecificData.SchemaConstructable; + +import java.nio.ByteBuffer; +import java.util.UUID; + +public interface Element extends IndexedRecord, SchemaConstructable { + Schema UUID_SCHEMA = + LogicalTypes.uuid().addToSchema(SchemaBuilder.builder().fixed("uuid").size(16)); + + static UUID toUuid(GenericData.Fixed fixed) { + ByteBuffer bb = ByteBuffer.wrap(fixed.bytes()); + long firstLong = bb.getLong(); + long secondLong = bb.getLong(); + return new UUID(firstLong, secondLong); + } + + static GenericData.Fixed toFixed(UUID uuid) { + ByteBuffer bb = ByteBuffer.allocate(16); + bb.putLong(uuid.getMostSignificantBits()); + bb.putLong(uuid.getLeastSignificantBits()); + return new GenericData.Fixed(UUID_SCHEMA, bb.array()); + } +} diff --git a/core/src/main/java/kafka/automq/table/events/Envelope.java b/core/src/main/java/kafka/automq/table/events/Envelope.java new file mode 100644 index 0000000000..02db9a3d59 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/Envelope.java @@ -0,0 +1,45 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +public class Envelope { + private final int partition; + private final long offset; + private final Event event; + + public Envelope(int partition, long offset, Event event) { + this.partition = partition; + this.offset = offset; + this.event = event; + } + + public int partition() { + return partition; + } + + public long offset() { + return offset; + } + + public Event event() { + return event; + } + +} diff --git a/core/src/main/java/kafka/automq/table/events/Errors.java b/core/src/main/java/kafka/automq/table/events/Errors.java new file mode 100644 index 0000000000..4bb168a5e7 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/Errors.java @@ -0,0 +1,27 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +public class Errors { + public static final int NONE = 0; + public static final int EPOCH_MISMATCH = 1; + public static final int MORE_DATA = 2; + +} diff --git a/core/src/main/java/kafka/automq/table/events/Event.java b/core/src/main/java/kafka/automq/table/events/Event.java new file mode 100644 index 0000000000..e52b53d5d0 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/Event.java @@ -0,0 +1,96 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; + +public class Event implements Element { + private long timestamp; + private EventType type; + private Payload payload; + private Schema avroSchema; + + // used by avro deserialize reflection + public Event(Schema avroSchema) { + this.avroSchema = avroSchema; + } + + public Event(long timestamp, EventType type, Payload payload) { + this.timestamp = timestamp; + this.type = type; + this.payload = payload; + avroSchema = SchemaBuilder.builder().record(Event.class.getName()) + .fields() + .name("timestamp").type().longType().noDefault() + .name("type").type().intType().noDefault() + .name("payload").type(payload.getSchema()).noDefault() + .endRecord(); + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.timestamp = (long) v; + return; + case 1: + this.type = EventType.fromId((Integer) v); + return; + case 2: + this.payload = (Payload) v; + return; + default: + // ignore the object, it must be from a newer version of the format + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return timestamp; + case 1: + return type.id(); + case 2: + return payload; + default: + throw new IllegalArgumentException("Unknown field index: " + i); + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + public long timestamp() { + return timestamp; + } + + public EventType type() { + return type; + } + + public T payload() { + //noinspection unchecked + return (T) payload; + } +} diff --git a/core/src/main/java/kafka/automq/table/events/EventType.java b/core/src/main/java/kafka/automq/table/events/EventType.java new file mode 100644 index 0000000000..67acc7ed39 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/EventType.java @@ -0,0 +1,46 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +public enum EventType { + COMMIT_REQUEST(0), + COMMIT_RESPONSE(1); + + private final int id; + + EventType(int id) { + this.id = id; + } + + public int id() { + return id; + } + + public static EventType fromId(int id) { + switch (id) { + case 0: + return COMMIT_REQUEST; + case 1: + return COMMIT_RESPONSE; + default: + throw new IllegalArgumentException("Unknown event type id: " + id); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/events/PartitionMetric.java b/core/src/main/java/kafka/automq/table/events/PartitionMetric.java new file mode 100644 index 0000000000..e0f9dfeb15 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/PartitionMetric.java @@ -0,0 +1,108 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; + +import java.util.Objects; + +public class PartitionMetric implements Element { + private int partition; + private long watermark; + private final Schema avroSchema; + + public static final Schema AVRO_SCHEMA = SchemaBuilder.builder() + .record(PartitionMetric.class.getName()) + .fields() + .name("partition").type().intType().noDefault() + .name("watermark").type().longType().noDefault() + .endRecord(); + + public PartitionMetric(Schema avroSchema) { + this.avroSchema = avroSchema; + } + + public PartitionMetric(int partition, long watermark) { + this.partition = partition; + this.watermark = watermark; + this.avroSchema = AVRO_SCHEMA; + } + + public int partition() { + return partition; + } + + public long watermark() { + return watermark; + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.partition = (int) v; + return; + case 1: + this.watermark = (long) v; + return; + default: + // ignore the object, it must be from a newer version of the format + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return partition; + case 1: + return watermark; + default: + throw new UnsupportedOperationException("Unknown field oridinal: " + i); + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + @Override + public String toString() { + return "PartitionMetrics{" + + "partition=" + partition + + ", watermark=" + watermark + + '}'; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + PartitionMetric metric = (PartitionMetric) o; + return partition == metric.partition && watermark == metric.watermark; + } + + @Override + public int hashCode() { + return Objects.hash(partition, watermark); + } +} diff --git a/core/src/main/java/kafka/automq/table/events/Payload.java b/core/src/main/java/kafka/automq/table/events/Payload.java new file mode 100644 index 0000000000..ed8c4eaca7 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/Payload.java @@ -0,0 +1,23 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +public interface Payload extends Element { +} diff --git a/core/src/main/java/kafka/automq/table/events/TopicMetric.java b/core/src/main/java/kafka/automq/table/events/TopicMetric.java new file mode 100644 index 0000000000..cd122fbec4 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/TopicMetric.java @@ -0,0 +1,94 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; + +import java.util.Objects; + +public class TopicMetric implements Element { + private long fieldCount; + private final Schema avroSchema; + public static final Schema AVRO_SCHEMA = SchemaBuilder.builder().record(TopicMetric.class.getName()).fields() + .name("fieldCount").type().longType().noDefault() + .endRecord(); + public static final TopicMetric NOOP = new TopicMetric(0); + + public TopicMetric(Schema avroSchema) { + this.avroSchema = avroSchema; + this.fieldCount = 0; + } + + public TopicMetric(long fieldCount) { + this.fieldCount = fieldCount; + this.avroSchema = AVRO_SCHEMA; + } + + public long fieldCount() { + return fieldCount; + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.fieldCount = (long) v; + return; + default: + // ignore the object, it must be from a newer version of the format + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return fieldCount; + default: + throw new UnsupportedOperationException("Unknown field oridinal: " + i); + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + TopicMetric metric = (TopicMetric) o; + return fieldCount == metric.fieldCount; + } + + @Override + public int hashCode() { + return Objects.hashCode(fieldCount); + } + + @Override + public String toString() { + return "TopicMetric{" + + "fieldCount=" + fieldCount + + '}'; + } +} diff --git a/core/src/main/java/kafka/automq/table/events/WorkerOffset.java b/core/src/main/java/kafka/automq/table/events/WorkerOffset.java new file mode 100644 index 0000000000..52aff1050a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/events/WorkerOffset.java @@ -0,0 +1,124 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; + +import java.util.Objects; + +public class WorkerOffset implements Element { + private int partition; + private int epoch; + private long offset; + private final Schema avroSchema; + + public static final Schema AVRO_SCHEMA = SchemaBuilder.builder() + .record(WorkerOffset.class.getName()) + .fields() + .name("partition") + .type().intType().noDefault() + .name("epoch") + .type().intType().noDefault() + .name("offset") + .type().longType().noDefault() + .endRecord(); + + public WorkerOffset(Schema avroSchema) { + this.avroSchema = avroSchema; + } + + public WorkerOffset(int partition, int epoch, long offset) { + this.partition = partition; + this.epoch = epoch; + this.offset = offset; + this.avroSchema = AVRO_SCHEMA; + } + + public int partition() { + return partition; + } + + public int epoch() { + return epoch; + } + + public long offset() { + return offset; + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.partition = (int) v; + return; + case 1: + this.epoch = (int) v; + return; + case 2: + this.offset = (long) v; + return; + default: + // ignore the object, it must be from a newer version of the format + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return partition; + case 1: + return epoch; + case 2: + return offset; + default: + throw new UnsupportedOperationException("Unknown field ordinal: " + i); + } + } + + @Override + public Schema getSchema() { + return avroSchema; + } + + @Override + public String toString() { + return "WorkerOffset{" + + "partition=" + partition + + ", epoch=" + epoch + + ", offset=" + offset + + '}'; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + WorkerOffset offset1 = (WorkerOffset) o; + return partition == offset1.partition && epoch == offset1.epoch && offset == offset1.offset; + } + + @Override + public int hashCode() { + return Objects.hash(partition, epoch, offset); + } +} diff --git a/core/src/main/java/kafka/automq/table/metric/FieldMetric.java b/core/src/main/java/kafka/automq/table/metric/FieldMetric.java new file mode 100644 index 0000000000..1f14578499 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/metric/FieldMetric.java @@ -0,0 +1,77 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.metric; + +import org.apache.avro.util.Utf8; + +import java.nio.ByteBuffer; + +public final class FieldMetric { + + private static final int STRING_BASE_COST = 3; // base cost for small strings + private static final int STRING_UNIT_BYTES = 32; // granularity for string scaling + private static final int STRING_UNIT_STEP = 1; // aggressive scaling for long strings + + private static final int BINARY_BASE_COST = 4; // small binary payloads slightly heavier than primitives + private static final int BINARY_UNIT_BYTES = 32; // granularity for binary buffers + private static final int BINARY_UNIT_STEP = 1; // scaling factor for binary payloads + + private FieldMetric() { + } + + public static int count(CharSequence value) { + if (value == null) { + return 0; + } + int lengthBytes = value instanceof Utf8 + ? ((Utf8) value).getByteLength() + : value.length(); + + if (lengthBytes <= STRING_UNIT_BYTES) { + return STRING_BASE_COST; + } + int segments = (lengthBytes + STRING_UNIT_BYTES - 1) / STRING_UNIT_BYTES; + return STRING_BASE_COST + (segments - 1) * STRING_UNIT_STEP; + } + + public static int count(ByteBuffer value) { + if (value == null) { + return 0; + } + int remaining = value.remaining(); + if (remaining <= BINARY_UNIT_BYTES) { + return BINARY_BASE_COST; + } + int segments = (remaining + BINARY_UNIT_BYTES - 1) / BINARY_UNIT_BYTES; + return BINARY_BASE_COST + (segments - 1) * BINARY_UNIT_STEP; + } + + public static int count(byte[] value) { + if (value == null) { + return 0; + } + int length = value.length; + if (length <= BINARY_UNIT_BYTES) { + return BINARY_BASE_COST; + } + int segments = (length + BINARY_UNIT_BYTES - 1) / BINARY_UNIT_BYTES; + return BINARY_BASE_COST + (segments - 1) * BINARY_UNIT_STEP; + } +} diff --git a/core/src/main/java/kafka/automq/table/metric/TableTopicMetricsManager.java b/core/src/main/java/kafka/automq/table/metric/TableTopicMetricsManager.java new file mode 100644 index 0000000000..200ec175b5 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/metric/TableTopicMetricsManager.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.metric; + +import com.automq.stream.s3.metrics.Metrics; +import com.automq.stream.s3.metrics.MetricsLevel; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; + +import java.time.Duration; +import java.util.concurrent.ExecutionException; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.Meter; + +public final class TableTopicMetricsManager { + private static final Cache TOPIC_ATTRIBUTE_CACHE = CacheBuilder.newBuilder() + .expireAfterAccess(Duration.ofMinutes(1)).build(); + private static final Metrics.LongGaugeBundle DELAY_GAUGES = Metrics.instance() + .longGauge("kafka_tabletopic_delay", "Table topic commit delay", "ms"); + private static final Metrics.DoubleGaugeBundle FIELDS_PER_SECOND_GAUGES = Metrics.instance() + .doubleGauge("kafka_tabletopic_fps", "Table topic fields per second", "fields/s"); + private static final Metrics.DoubleGaugeBundle EVENT_LOOP_BUSY_GAUGES = Metrics.instance() + .doubleGauge("kafka_tableworker_eventloop_busy_ratio", "Table worker event loop busy ratio", "%"); + + private TableTopicMetricsManager() { + } + + public static void initMetrics(Meter meter) { + // Metrics instruments are registered via Metrics.instance(); no additional setup required. + } + + public static Metrics.LongGaugeBundle.LongGauge registerDelay(String topic) { + return DELAY_GAUGES.register(MetricsLevel.INFO, getTopicAttribute(topic)); + } + + public static Metrics.DoubleGaugeBundle.DoubleGauge registerFieldsPerSecond(String topic) { + return FIELDS_PER_SECOND_GAUGES.register(MetricsLevel.INFO, getTopicAttribute(topic)); + } + + public static Metrics.DoubleGaugeBundle.DoubleGauge registerEventLoopBusy(String loop) { + return EVENT_LOOP_BUSY_GAUGES.register(MetricsLevel.INFO, getLoopAttribute(loop)); + } + + private static Attributes getTopicAttribute(String topic) { + try { + return TOPIC_ATTRIBUTE_CACHE.get(topic, () -> Attributes.of(AttributeKey.stringKey("topic"), topic)); + } catch (ExecutionException e) { + throw new RuntimeException(e); + } + } + + private static Attributes getLoopAttribute(String loop) { + return Attributes.of(AttributeKey.stringKey("event_loop"), loop); + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/AvroTestCase.java b/core/src/main/java/kafka/automq/table/perf/AvroTestCase.java new file mode 100644 index 0000000000..dd4a365d47 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/AvroTestCase.java @@ -0,0 +1,155 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.convert.AvroRegistryConverter; + +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.Encoder; +import org.apache.avro.io.EncoderFactory; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +class AvroTestCase extends PerfTestCase { + private final Schema schema; + private final PayloadManager payloadManager; + + public AvroTestCase(DataType dataType, int fieldCount, int payloadCount, PerfConfig config) { + super(dataType, "avro", fieldCount, payloadCount, config); + this.schema = createSchema(dataType, fieldCount); + this.payloadManager = new PayloadManager(() -> generatePayloadWithSchema(schema), payloadCount); + } + + @Override + PayloadManager getPayloadManager() { + return payloadManager; + } + + @Override + protected byte[] generatePayload() { + return generatePayloadWithSchema(schema); + } + + private byte[] generatePayloadWithSchema(Schema schemaToUse) { + GenericRecord record = new GenericData.Record(schemaToUse); + fillRecord(record, dataType, fieldCount); + return createAvroValue(record); + } + + @Override + protected Converter createConverter() { + StaticAvroDeserializer deserializer = new StaticAvroDeserializer(schema); + return new AvroRegistryConverter(deserializer, null); + } + + private Schema createSchema(DataType dataType, int fieldCount) { + SchemaBuilder.FieldAssembler assembler = SchemaBuilder.builder() + .record("TestRecord") + .fields(); + + for (int i = 0; i < fieldCount; i++) { + String fieldName = "f" + i; + assembler = addFieldToSchema(assembler, fieldName, dataType); + } + + return assembler.endRecord(); + } + + private SchemaBuilder.FieldAssembler addFieldToSchema(SchemaBuilder.FieldAssembler assembler, String fieldName, DataType dataType) { + return switch (dataType) { + case BOOLEAN -> assembler.name(fieldName).type().booleanType().noDefault(); + case INT -> assembler.name(fieldName).type().intType().noDefault(); + case LONG -> assembler.name(fieldName).type().longType().noDefault(); + case DOUBLE -> assembler.name(fieldName).type().doubleType().noDefault(); + case TIMESTAMP -> assembler.name(fieldName).type(LogicalTypes.timestampMillis() + .addToSchema(Schema.create(Schema.Type.LONG))).withDefault(0); + case STRING -> assembler.name(fieldName).type().stringType().noDefault(); + case BINARY -> assembler.name(fieldName).type().bytesType().noDefault(); + case NESTED -> { + Schema nestedSchema = SchemaBuilder.builder().record("nested").fields() + .name("nf1").type().booleanType().noDefault() + .endRecord(); + yield assembler.name(fieldName).type(nestedSchema).noDefault(); + } + case ARRAY -> assembler.name(fieldName).type().array().items(Schema.create(Schema.Type.BOOLEAN)).noDefault(); + }; + } + + private void fillRecord(GenericRecord record, DataType dataType, int fieldCount) { + for (int i = 0; i < fieldCount; i++) { + String fieldName = "f" + i; + Object value = generateFieldValue(dataType); + record.put(fieldName, value); + } + } + + private Object generateFieldValue(DataType dataType) { + return switch (dataType) { + case BOOLEAN, INT, LONG, DOUBLE, TIMESTAMP, STRING -> dataType.generateValue(); + case BINARY -> { + ByteBuffer buffer = (ByteBuffer) dataType.generateValue(); + yield buffer; + } + case NESTED -> { + // Create nested schema inline + Schema nestedSchema = SchemaBuilder.builder().record("nested").fields() + .name("nf1").type().booleanType().noDefault() + .endRecord(); + GenericRecord nestedRecord = new GenericData.Record(nestedSchema); + nestedRecord.put("nf1", DataType.BOOLEAN.generateValue()); + yield nestedRecord; + } + case ARRAY -> List.of(DataType.BOOLEAN.generateValue()); + }; + } + + private static byte[] createAvroValue(GenericRecord record) { + try { + DatumWriter datumWriter = new GenericDatumWriter<>(record.getSchema()); + ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + Encoder encoder = EncoderFactory.get().binaryEncoder(outputStream, null); + datumWriter.write(record, encoder); + encoder.flush(); + byte[] avroBytes = outputStream.toByteArray(); + + ByteBuf buf = Unpooled.buffer(1 + 4 + avroBytes.length); + buf.writeByte((byte) 0x0); + buf.writeInt(0); + buf.writeBytes(avroBytes); + + return buf.array(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/BenchmarkResult.java b/core/src/main/java/kafka/automq/table/perf/BenchmarkResult.java new file mode 100644 index 0000000000..a97d41d730 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/BenchmarkResult.java @@ -0,0 +1,112 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +public class BenchmarkResult { + private final String formatName; + private final String dataTypeName; + private final long durationNs; + private final long recordsProcessed; + private final long fieldCount; + private final String errorMessage; + + private BenchmarkResult(String formatName, String dataTypeName, long durationNs, long recordsProcessed, long fieldCount, + String errorMessage) { + this.formatName = formatName; + this.dataTypeName = dataTypeName; + this.durationNs = durationNs; + this.recordsProcessed = recordsProcessed; + this.fieldCount = fieldCount; + this.errorMessage = errorMessage; + } + + public static BenchmarkResult success(String formatName, String dataTypeName, long durationNs, + long recordsProcessed, long fieldCount) { + return new BenchmarkResult(formatName, dataTypeName, durationNs, recordsProcessed, fieldCount, null); + } + + public static BenchmarkResult failure(String formatName, String dataTypeName, String errorMessage) { + return new BenchmarkResult(formatName, dataTypeName, 0, 0, 0, errorMessage); + } + + public String getFormatName() { + return formatName; + } + + public String getDataTypeName() { + return dataTypeName; + } + + public long getDurationNs() { + return durationNs; + } + + public long getDurationMs() { + return durationNs / 1_000_000L; + } + + public long getRecordsProcessed() { + return recordsProcessed; + } + + public long getFieldCount() { + return fieldCount; + } + + public String getErrorMessage() { + return errorMessage; + } + + public boolean isSuccess() { + return errorMessage == null; + } + + public long getThroughput() { + long durationMs = getDurationMs(); + if (durationMs == 0) { + return 0; + } + return (recordsProcessed * 1000L) / durationMs; + } + + public double getNsPerField() { + if (fieldCount == 0) { + return 0.0d; + } + return (double) durationNs / (double) fieldCount; + } + + public double getNsPerRecord() { + if (recordsProcessed == 0) { + return 0.0d; + } + return (double) durationNs / (double) recordsProcessed; + } + + @Override + public String toString() { + if (isSuccess()) { + return String.format("%s %s: %d ms, %d records, fieldCount=%d, ns/field=%.2f", + formatName, dataTypeName, getDurationMs(), recordsProcessed, fieldCount, getNsPerField()); + } else { + return String.format("%s %s: FAILED - %s", formatName, dataTypeName, errorMessage); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/BenchmarkWorkerConfig.java b/core/src/main/java/kafka/automq/table/perf/BenchmarkWorkerConfig.java new file mode 100644 index 0000000000..a0e1529a7a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/BenchmarkWorkerConfig.java @@ -0,0 +1,112 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import kafka.automq.table.worker.WorkerConfig; + +import org.apache.kafka.server.record.ErrorsTolerance; +import org.apache.kafka.server.record.TableTopicConvertType; +import org.apache.kafka.server.record.TableTopicSchemaType; +import org.apache.kafka.server.record.TableTopicTransformType; + +import java.util.Collections; +import java.util.List; + +class BenchmarkWorkerConfig extends WorkerConfig { + + public BenchmarkWorkerConfig() { + super(); + } + + @Override + public String namespace() { + return "test"; + } + + @Override + public TableTopicSchemaType schemaType() { + return TableTopicSchemaType.NONE; + } + + @Override + public TableTopicConvertType valueConvertType() { + return TableTopicConvertType.BY_SCHEMA_ID; + } + + @Override + public TableTopicConvertType keyConvertType() { + return TableTopicConvertType.STRING; + } + + @Override + public TableTopicTransformType transformType() { + return TableTopicTransformType.FLATTEN; + } + + @Override + public String valueSubject() { + return null; + } + + @Override + public String valueMessageFullName() { + return null; + } + + @Override + public String keySubject() { + return null; + } + + @Override + public String keyMessageFullName() { + return null; + } + + @Override + public List idColumns() { + return Collections.emptyList(); + } + + @Override + public String partitionByConfig() { + return null; + } + + @Override + public List partitionBy() { + return Collections.emptyList(); + } + + @Override + public boolean upsertEnable() { + return false; + } + + @Override + public ErrorsTolerance errorsTolerance() { + return ErrorsTolerance.ALL; + } + + @Override + public String cdcField() { + return null; + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/DataType.java b/core/src/main/java/kafka/automq/table/perf/DataType.java new file mode 100644 index 0000000000..a167b51e73 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/DataType.java @@ -0,0 +1,81 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import org.apache.commons.lang3.RandomStringUtils; + +import java.nio.ByteBuffer; +import java.util.concurrent.ThreadLocalRandom; +import java.util.function.Supplier; + +public enum DataType { + BOOLEAN("boolean", () -> ThreadLocalRandom.current().nextBoolean()), + INT("int", () -> ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)), + LONG("long", () -> ThreadLocalRandom.current().nextLong(Long.MAX_VALUE)), + DOUBLE("double", () -> ThreadLocalRandom.current().nextDouble(Long.MAX_VALUE)), + TIMESTAMP("timestamp", () -> ThreadLocalRandom.current().nextLong(Long.MAX_VALUE)), + STRING("string", () -> RandomStringUtils.randomAlphabetic(32)), + BINARY("binary", () -> { + byte[] bytes = new byte[32]; + ThreadLocalRandom.current().nextBytes(bytes); + return ByteBuffer.wrap(bytes); + }), + NESTED("nested", null), + ARRAY("array", null); + + private final String name; + // Supplier holds runtime-only generators (often non-serializable lambdas). Enum + // instances are serialized by name, not by fields; mark this field transient to + // avoid accidental Java serialization of the supplier and silence static analyzers. + private final transient Supplier valueGenerator; + + DataType(String name, Supplier valueGenerator) { + this.name = name; + this.valueGenerator = valueGenerator; + } + + public String getName() { + return name; + } + + public Object generateValue() { + if (valueGenerator == null) { + throw new UnsupportedOperationException("Complex type " + name + " requires specific generator"); + } + return valueGenerator.get(); + } + + public static DataType fromString(String name) { + for (DataType type : values()) { + if (type.name.equals(name)) { + return type; + } + } + return null; + } + + public PerfTestCase createAvroTestCase(int fieldCount, int payloadCount, PerfConfig config) { + return new AvroTestCase(this, fieldCount, payloadCount, config); + } + + public PerfTestCase createProtobufTestCase(int fieldCount, int payloadCount, PerfConfig config) { + return new ProtobufTestCase(this, fieldCount, payloadCount, config); + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/Deserializers.java b/core/src/main/java/kafka/automq/table/perf/Deserializers.java new file mode 100644 index 0000000000..19c6276630 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/Deserializers.java @@ -0,0 +1,71 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import org.apache.kafka.common.serialization.Deserializer; + +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Message; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.DatumReader; +import org.apache.avro.io.DecoderFactory; + +import java.io.ByteArrayInputStream; +import java.io.IOException; + +class StaticAvroDeserializer implements Deserializer { + final Schema schema; + final DatumReader reader; + final DecoderFactory decoderFactory = DecoderFactory.get(); + + public StaticAvroDeserializer(Schema schema) { + this.schema = schema; + this.reader = new GenericDatumReader<>(schema); + } + + public Object deserialize(String topic, byte[] data) { + try { + return this.reader.read(null, decoderFactory.binaryDecoder(data, 5, data.length - 5, null)); + } catch (IOException e) { + throw new RuntimeException(e); + } + } +} + +class StaticProtobufDeserializer implements Deserializer { + final Descriptors.Descriptor descriptor; + + public StaticProtobufDeserializer(Descriptors.Descriptor descriptor) { + this.descriptor = descriptor; + } + + @Override + public Message deserialize(String s, byte[] bytes) { + try { + return DynamicMessage.parseFrom(descriptor, new ByteArrayInputStream(bytes, 5, bytes.length - 5)); + } catch (IOException e) { + throw new RuntimeException(e); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/FieldProfileRunner.java b/core/src/main/java/kafka/automq/table/perf/FieldProfileRunner.java new file mode 100644 index 0000000000..5a6d02cb2d --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/FieldProfileRunner.java @@ -0,0 +1,536 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.convert.AvroRegistryConverter; + +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.Encoder; +import org.apache.avro.io.EncoderFactory; +import org.apache.commons.lang3.RandomStringUtils; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ThreadLocalRandom; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +/** + * Runs a suite of single-field conversion benchmarks to compare end-to-end CPU cost per field type/shape. + */ +public final class FieldProfileRunner { + + private static final List DEFAULT_SCENARIOS = List.of( + FieldScenario.primitive("boolean", DataType.BOOLEAN, Schema.create(Schema.Type.BOOLEAN)), + FieldScenario.primitive("int", DataType.INT, Schema.create(Schema.Type.INT)), + FieldScenario.primitive("long", DataType.LONG, Schema.create(Schema.Type.LONG)), + FieldScenario.primitive("double", DataType.DOUBLE, Schema.create(Schema.Type.DOUBLE)), + FieldScenario.logical("timestamp-millis", DataType.TIMESTAMP, + LogicalTypes.timestampMillis().addToSchema(Schema.create(Schema.Type.LONG))), + FieldScenario.string("string-8", 8), + FieldScenario.string("string-32", 32), + FieldScenario.string("string-128", 128), + FieldScenario.string("string-256", 256), + FieldScenario.string("string-512", 512), + FieldScenario.string("string-2048", 2048), + FieldScenario.string("string-4096", 4096), + FieldScenario.binary("binary-8", 8), + FieldScenario.binary("binary-32", 32), + FieldScenario.binary("binary-128", 128), + FieldScenario.binary("binary-256", 256), + FieldScenario.binary("binary-512", 512), + FieldScenario.binary("binary-2048", 2048), + FieldScenario.fixed("fixed-16", 16), + FieldScenario.fixed("fixed-128", 128), + FieldScenario.decimal("decimal-38-4", 38, 4), + FieldScenario.uuid("uuid"), + FieldScenario.list("list-int-16", Schema.create(Schema.Type.INT), 16), + FieldScenario.map("map-string-long-16", Schema.create(Schema.Type.LONG), 16), + FieldScenario.struct("struct-nested", List.of( + Schema.create(Schema.Type.INT), + Schema.create(Schema.Type.LONG), + Schema.create(Schema.Type.STRING))) + ); + + private FieldProfileRunner() { + } + + public static void main(String[] args) { + ProfileConfig config = ProfileConfig.fromEnv(); + List scenarios = config.resolvedScenarios(DEFAULT_SCENARIOS); + Map results = new LinkedHashMap<>(); + + System.out.printf("Field profile records=%d payloads=%d batchBytes=%d rounds=%d%n", + config.recordsCount, config.payloadCount, config.batchSizeBytes, config.rounds); + + for (FieldScenario scenario : scenarios) { + SingleFieldAvroTestCase testCase = new SingleFieldAvroTestCase( + scenario, + config.recordsCount, + config.payloadCount, + config.batchSizeBytes); + + System.out.printf("Scenario: %s (%s)%n", scenario.name, scenario.describe()); + + long totalDurationNs = 0L; + long totalFieldCount = 0L; + long totalRecords = 0L; + boolean success = true; + + for (int round = 1; round <= config.rounds; round++) { + BenchmarkResult result = testCase.runBenchmark(config.recordsCount); + if (!result.isSuccess()) { + System.out.printf(" round %d failed: %s%n", round, result.getErrorMessage()); + success = false; + break; + } + + totalDurationNs += result.getDurationNs(); + totalFieldCount += result.getFieldCount(); + totalRecords += result.getRecordsProcessed(); + + System.out.printf(Locale.ROOT, + " round %d: %d ms, fieldCount=%d, ns/field=%.2f, ns/record=%.2f%n", + round, + result.getDurationMs(), + result.getFieldCount(), + result.getNsPerField(), + result.getNsPerRecord()); + } + + if (success && totalRecords > 0) { + BenchmarkResult aggregated = BenchmarkResult.success( + testCase.formatName, + scenario.name, + totalDurationNs, + totalRecords, + totalFieldCount); + results.put(scenario.name, aggregated); + System.out.printf(Locale.ROOT, + " aggregated: ns/field=%.2f, ns/record=%.2f (over %d rounds)%n", + aggregated.getNsPerField(), + aggregated.getNsPerRecord(), + config.rounds); + } else { + results.put(scenario.name, BenchmarkResult.failure(testCase.formatName, scenario.name, + "round failed")); + } + } + + summarize(results); + } + + private static void summarize(Map results) { + BenchmarkResult primitiveBaseline = results.get("int"); + double baselineNs = primitiveBaseline != null ? primitiveBaseline.getNsPerField() : 0.0d; + + System.out.println(); + System.out.println("=== Summary ==="); + System.out.printf(Locale.ROOT, "%-24s %-12s %-12s %-12s %-12s%n", + "scenario", "ns/field", "ns/record", "fieldCount", "rel"); + + results.forEach((name, result) -> { + if (!result.isSuccess()) { + System.out.printf("%-24s %-12s %-12s %-12s %-12s%n", + name, "FAIL", "-", "-", "-"); + return; + } + double rel = baselineNs == 0.0d ? 0.0d : result.getNsPerField() / baselineNs; + System.out.printf(Locale.ROOT, "%-24s %-12.2f %-12.2f %-12d %-12.2f%n", + name, + result.getNsPerField(), + result.getNsPerRecord(), + result.getFieldCount(), + rel); + }); + } + + private static final class ProfileConfig { + final long recordsCount; + final int payloadCount; + final int batchSizeBytes; + final Set scenarioFilter; + final int rounds; + + private ProfileConfig(long recordsCount, int payloadCount, int batchSizeBytes, Set scenarioFilter, int rounds) { + this.recordsCount = recordsCount; + this.payloadCount = payloadCount; + this.batchSizeBytes = batchSizeBytes; + this.scenarioFilter = scenarioFilter; + this.rounds = rounds; + } + + static ProfileConfig fromEnv() { + long records = parseLongEnv("FIELD_PROFILE_RECORDS", 2_000_000L); + int payloads = (int) parseLongEnv("FIELD_PROFILE_PAYLOADS", 1_024L); + int batchSize = (int) parseLongEnv("FIELD_PROFILE_BATCH_BYTES", 32 * 1024 * 1024L); + Set filter = EnvParsers.parseCsvToSet(System.getenv("FIELD_PROFILE_SCENARIOS")); + int rounds = (int) parseLongEnv("FIELD_PROFILE_ROUNDS", 3L); + return new ProfileConfig(records, payloads, batchSize, filter, Math.max(rounds, 1)); + } + + List resolvedScenarios(List defaults) { + if (scenarioFilter == null || scenarioFilter.isEmpty()) { + return defaults; + } + List filtered = new ArrayList<>(); + for (FieldScenario scenario : defaults) { + if (scenarioFilter.contains(scenario.name)) { + filtered.add(scenario); + } + } + return filtered.isEmpty() ? defaults : filtered; + } + + private static long parseLongEnv(String key, long defaultValue) { + String value = System.getenv(key); + if (value == null || value.isEmpty()) { + return defaultValue; + } + try { + return Long.parseLong(value); + } catch (NumberFormatException e) { + return defaultValue; + } + } + } + + private static final class SingleFieldAvroTestCase extends PerfTestCase { + private final FieldScenario scenario; + private final Schema recordSchema; + private final PayloadManager payloadManager; + private final StaticAvroDeserializer deserializer; + + SingleFieldAvroTestCase(FieldScenario scenario, long recordsCount, int payloadCount, int batchSizeBytes) { + super(scenario.baseDataType, "avro", scenario.fieldCount, payloadCount, + new InlinePerfConfig(recordsCount, scenario.fieldCount, payloadCount, batchSizeBytes)); + this.scenario = scenario; + this.recordSchema = scenario.buildRecordSchema(); + this.deserializer = new StaticAvroDeserializer(recordSchema); + this.payloadManager = new PayloadManager(this::generatePayload, payloadCount); + } + + @Override + PayloadManager getPayloadManager() { + return payloadManager; + } + + @Override + protected byte[] generatePayload() { + GenericRecord record = new GenericData.Record(recordSchema); + scenario.populate(record); + return encodeWithHeader(recordSchema, record); + } + + @Override + protected Converter createConverter() { + return new AvroRegistryConverter(deserializer, null); + } + + private static byte[] encodeWithHeader(Schema schema, GenericRecord record) { + try { + DatumWriter datumWriter = new GenericDatumWriter<>(schema); + ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + Encoder encoder = EncoderFactory.get().binaryEncoder(outputStream, null); + datumWriter.write(record, encoder); + encoder.flush(); + byte[] avroBytes = outputStream.toByteArray(); + + ByteBuf buffer = Unpooled.buffer(1 + 4 + avroBytes.length); + buffer.writeByte((byte) 0x0); + buffer.writeInt(0); + buffer.writeBytes(avroBytes); + return buffer.array(); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + } + + private static final class InlinePerfConfig extends PerfConfig { + private final long recordsCount; + private final int fieldCount; + private final int payloadsCount; + private final int batchSizeBytes; + + InlinePerfConfig(long recordsCount, int fieldCount, int payloadsCount, int batchSizeBytes) { + this.recordsCount = recordsCount; + this.fieldCount = fieldCount; + this.payloadsCount = payloadsCount; + this.batchSizeBytes = batchSizeBytes; + } + + @Override + public long getRecordsCount() { + return recordsCount; + } + + @Override + public int getFieldCountPerRecord() { + return fieldCount; + } + + @Override + public int getPayloadsCount() { + return payloadsCount; + } + + @Override + public int getBatchSizeBytes() { + return batchSizeBytes; + } + } + + private static final class FieldScenario { + final String name; + final DataType baseDataType; + final int fieldCount; + final Schema fieldSchema; + final ScenarioKind kind; + final int size; + final int decimalPrecision; + final int decimalScale; + final List structFieldSchemas; + + private FieldScenario(String name, DataType baseDataType, int fieldCount, Schema fieldSchema, ScenarioKind kind, int size) { + this(name, baseDataType, fieldCount, fieldSchema, kind, size, 0, 0, List.of()); + } + + private FieldScenario(String name, DataType baseDataType, int fieldCount, Schema fieldSchema, ScenarioKind kind, + int size, int decimalPrecision, int decimalScale, List structFieldSchemas) { + this.name = name; + this.baseDataType = baseDataType; + this.fieldCount = fieldCount; + this.fieldSchema = fieldSchema; + this.kind = kind; + this.size = size; + this.decimalPrecision = decimalPrecision; + this.decimalScale = decimalScale; + this.structFieldSchemas = structFieldSchemas; + } + + static FieldScenario primitive(String name, DataType dataType, Schema schema) { + return new FieldScenario(name, dataType, 1, schema, ScenarioKind.PRIMITIVE, 0); + } + + static FieldScenario logical(String name, DataType dataType, Schema schema) { + return new FieldScenario(name, dataType, 1, schema, ScenarioKind.PRIMITIVE, 0); + } + + static FieldScenario string(String name, int length) { + return new FieldScenario(name, DataType.STRING, 1, Schema.create(Schema.Type.STRING), ScenarioKind.STRING, length); + } + + static FieldScenario binary(String name, int length) { + return new FieldScenario(name, DataType.BINARY, 1, Schema.create(Schema.Type.BYTES), ScenarioKind.BINARY, length); + } + + static FieldScenario fixed(String name, int length) { + String fixedName = name.replace('-', '_') + "_fixed"; + return new FieldScenario(name, DataType.BINARY, 1, Schema.createFixed(fixedName, null, null, length), ScenarioKind.FIXED, length); + } + + static FieldScenario list(String name, Schema elementSchema, int elementCount) { + Schema listSchema = Schema.createArray(elementSchema); + return new FieldScenario(name, DataType.ARRAY, 1, listSchema, ScenarioKind.LIST, elementCount); + } + + static FieldScenario map(String name, Schema valueSchema, int entryCount) { + Schema mapSchema = Schema.createMap(valueSchema); + return new FieldScenario(name, DataType.NESTED, 1, mapSchema, ScenarioKind.MAP, entryCount); + } + + static FieldScenario decimal(String name, int precision, int scale) { + Schema decimalSchema = LogicalTypes.decimal(precision, scale).addToSchema(Schema.create(Schema.Type.BYTES)); + return new FieldScenario(name, DataType.BINARY, 1, decimalSchema, ScenarioKind.DECIMAL, 0, precision, scale, List.of()); + } + + static FieldScenario uuid(String name) { + Schema uuidSchema = LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING)); + return new FieldScenario(name, DataType.STRING, 1, uuidSchema, ScenarioKind.UUID, 0); + } + + static FieldScenario struct(String name, List fieldSchemas) { + SchemaBuilder.FieldAssembler assembler = SchemaBuilder.record(name.replace('-', '_') + "_struct") + .fields(); + for (int i = 0; i < fieldSchemas.size(); i++) { + assembler.name("sf" + i).type(fieldSchemas.get(i)).noDefault(); + } + Schema structSchema = assembler.endRecord(); + return new FieldScenario(name, DataType.NESTED, 1, structSchema, ScenarioKind.STRUCT, fieldSchemas.size(), 0, 0, fieldSchemas); + } + + Schema buildRecordSchema() { + SchemaBuilder.FieldAssembler assembler = SchemaBuilder.record(name.replace('-', '_')) + .namespace("kafka.automq.table.perf") + .fields(); + assembler.name("f0").type(fieldSchema).noDefault(); + return assembler.endRecord(); + } + + void populate(GenericRecord record) { + record.put("f0", value()); + } + + Object value() { + ThreadLocalRandom random = ThreadLocalRandom.current(); + return switch (kind) { + case PRIMITIVE -> primitiveValue(random); + case STRING -> new org.apache.avro.util.Utf8(randomString(size)); + case BINARY -> ByteBuffer.wrap(randomBytes(size)); + case FIXED -> { + GenericData.Fixed fixed = new GenericData.Fixed(fieldSchema); + fixed.bytes(randomBytes(size)); + yield fixed; + } + case DECIMAL -> decimalValue(); + case UUID -> new org.apache.avro.util.Utf8(UUID.randomUUID().toString()); + case LIST -> { + List list = new ArrayList<>(size); + for (int i = 0; i < size; i++) { + list.add(ThreadLocalRandom.current().nextInt(2) == 0 ? 0 : 1); + } + yield list; + } + case MAP -> { + Map map = new LinkedHashMap<>(size); + for (int i = 0; i < size; i++) { + map.put("k" + i, random.nextLong()); + } + yield map; + } + case STRUCT -> { + GenericRecord nested = new GenericData.Record(fieldSchema); + for (Schema.Field field : fieldSchema.getFields()) { + nested.put(field.name(), randomValueForSchema(field.schema())); + } + yield nested; + } + }; + } + + String describe() { + return switch (kind) { + case PRIMITIVE -> fieldSchema.getType().getName(); + case STRING -> "string(len=" + size + ")"; + case BINARY -> "bytes(len=" + size + ")"; + case FIXED -> "fixed(len=" + size + ")"; + case DECIMAL -> "decimal(precision=" + decimalPrecision + ",scale=" + decimalScale + ")"; + case UUID -> "uuid"; + case LIST -> "list[size=" + size + "]"; + case MAP -> "map[size=" + size + "]"; + case STRUCT -> "struct(fields=" + size + ")"; + }; + } + + private Object primitiveValue(ThreadLocalRandom random) { + return switch (fieldSchema.getType()) { + case BOOLEAN -> random.nextBoolean(); + case INT -> random.nextInt(); + case LONG -> random.nextLong(); + case DOUBLE -> random.nextDouble(); + default -> random.nextLong(); + }; + } + + private static byte[] randomBytes(int size) { + byte[] bytes = new byte[size]; + ThreadLocalRandom.current().nextBytes(bytes); + return bytes; + } + + private static String randomString(int size) { + return RandomStringUtils.randomAlphanumeric(size); + } + + private Object decimalValue() { + ThreadLocalRandom random = ThreadLocalRandom.current(); + BigDecimal value = BigDecimal.valueOf(random.nextDouble()) + .multiply(BigDecimal.TEN.pow(decimalPrecision - Math.max(decimalScale, 1))) + .setScale(decimalScale, RoundingMode.HALF_UP); + byte[] bytes = value.unscaledValue().toByteArray(); + return ByteBuffer.wrap(bytes); + } + + private Object randomValueForSchema(Schema schema) { + ThreadLocalRandom random = ThreadLocalRandom.current(); + return switch (schema.getType()) { + case BOOLEAN -> random.nextBoolean(); + case INT -> random.nextInt(); + case LONG -> random.nextLong(); + case DOUBLE -> random.nextDouble(); + case STRING -> new org.apache.avro.util.Utf8(randomString(8)); + case BYTES -> ByteBuffer.wrap(randomBytes(16)); + default -> null; + }; + } + } + + private enum ScenarioKind { + PRIMITIVE, + STRING, + BINARY, + FIXED, + DECIMAL, + UUID, + LIST, + MAP, + STRUCT + } + + private static final class EnvParsers { + private EnvParsers() { + } + + static Set parseCsvToSet(String csv) { + if (csv == null || csv.isEmpty()) { + return null; + } + String[] parts = csv.split(","); + Set set = new LinkedHashSet<>(parts.length); + for (String part : parts) { + String trimmed = part.trim(); + if (!trimmed.isEmpty()) { + set.add(trimmed); + } + } + return set; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/FieldsPerformanceTest.java b/core/src/main/java/kafka/automq/table/perf/FieldsPerformanceTest.java new file mode 100644 index 0000000000..24b8679b61 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/FieldsPerformanceTest.java @@ -0,0 +1,99 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import org.apache.commons.lang3.tuple.Pair; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +public class FieldsPerformanceTest { + + public static void main(String[] args) { + PerfConfig config = new PerfConfig(); + + Map> results = new HashMap<>(); + results.put("avro", new ArrayList<>()); + results.put("proto", new ArrayList<>()); + + System.out.printf("Starting performance tests with %d records per test%n", config.getRecordsCount()); + System.out.printf("Enabled data types: %s%n", config.getEnabledDataTypes()); + System.out.printf("Enabled formats: %s%n", config.getEnabledFormats()); + + for (DataType dataType : config.getEnabledDataTypes()) { + for (SerializationFormat format : config.getEnabledFormats()) { + + PerfTestCase testCase = createTestCase(dataType, format, config); + + try { + PerfTestCase.clearInMemoryFiles(); + + System.out.printf("Running benchmark: %s %s%n", format.getName(), dataType.getName()); + BenchmarkResult result = testCase.runBenchmark(config.getRecordsCount()); + + if (result.isSuccess()) { + results.get(format.getName()).add(result); + System.out.printf( + "Completed: %s %s - %d ms, fieldCount=%d, ns/field=%s%n", + format.getName(), + dataType.getName(), + result.getDurationMs(), + result.getFieldCount(), + String.format(java.util.Locale.ROOT, "%.2f", result.getNsPerField())); + } else { + System.err.printf("Failed: %s %s - %s%n", + format.getName(), dataType.getName(), result.getErrorMessage()); + } + + } catch (Exception e) { + System.err.printf("Failed: %s %s - %s%n", format.getName(), dataType.getName(), e.getMessage()); + e.printStackTrace(System.err); + } + } + } + + // Output results in the same format as original + results.forEach((format, formatResults) -> { + System.out.printf("type: %s%n", format); + List> durations = formatResults.stream() + .map(r -> Pair.of(r.getDataTypeName(), r.getDurationMs())) + .collect(Collectors.toList()); + System.out.printf("task cost: %s%n", durations); + formatResults.forEach(r -> System.out.printf( + "detail: %s %s -> records=%d, fieldCount=%d, ns/field=%s, ns/record=%s%n", + format, + r.getDataTypeName(), + r.getRecordsProcessed(), + r.getFieldCount(), + String.format(java.util.Locale.ROOT, "%.2f", r.getNsPerField()), + String.format(java.util.Locale.ROOT, "%.2f", r.getNsPerRecord()))); + }); + } + + private static PerfTestCase createTestCase(DataType dataType, SerializationFormat format, PerfConfig config) { + return switch (format) { + case AVRO -> dataType.createAvroTestCase(config.getFieldCountPerRecord(), config.getPayloadsCount(), config); + case PROTOBUF -> dataType.createProtobufTestCase(config.getFieldCountPerRecord(), config.getPayloadsCount(), config); + }; + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/PayloadManager.java b/core/src/main/java/kafka/automq/table/perf/PayloadManager.java new file mode 100644 index 0000000000..012054b019 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/PayloadManager.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import java.util.function.Supplier; + +public class PayloadManager { + private final byte[][] payloadPool; + private final int poolSize; + private int currentIndex = 0; + + public PayloadManager(Supplier payloadGenerator, int poolSize) { + this.poolSize = Math.min(poolSize, 10000); // Limit max pool size + this.payloadPool = new byte[this.poolSize][]; + + // Pre-generate all payloads + for (int i = 0; i < this.poolSize; i++) { + this.payloadPool[i] = payloadGenerator.get(); + } + } + + public byte[] nextPayload() { + byte[] payload = payloadPool[currentIndex]; + currentIndex = (currentIndex + 1) % poolSize; + return payload; + } + + public void reset() { + currentIndex = 0; + } + + public int getPoolSize() { + return poolSize; + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/PerfConfig.java b/core/src/main/java/kafka/automq/table/perf/PerfConfig.java new file mode 100644 index 0000000000..64f6c7708a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/PerfConfig.java @@ -0,0 +1,103 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import com.automq.stream.utils.Systems; + +import org.apache.commons.lang3.StringUtils; + +import java.util.Arrays; +import java.util.EnumSet; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +public class PerfConfig { + private final long recordsCount; + private final int fieldCountPerRecord; + private final int payloadsCount; + private final int batchSizeBytes; + private final Set enabledDataTypes; + private final Set enabledFormats; + + public PerfConfig() { + this.recordsCount = Systems.getEnvLong("RECORDS_COUNT", 10_000_000L); + this.fieldCountPerRecord = parseIntEnv("FIELD_COUNT_PER_RECORD", 32); + this.payloadsCount = parseIntEnv("PAYLOADS_COUNT", 1000); + this.batchSizeBytes = parseIntEnv("BATCH_SIZE_BYTES", 32 * 1024 * 1024); + this.enabledDataTypes = parseDataTypes(System.getenv("TASKS")); + this.enabledFormats = parseFormats(System.getenv("FORMAT_TYPES")); + } + + private int parseIntEnv(String envName, int defaultValue) { + String value = System.getenv(envName); + if (StringUtils.isBlank(value)) { + return defaultValue; + } + try { + return Integer.parseInt(value); + } catch (NumberFormatException e) { + return defaultValue; + } + } + + private Set parseDataTypes(String tasksStr) { + if (StringUtils.isBlank(tasksStr)) { + return EnumSet.allOf(DataType.class); + } + return Arrays.stream(tasksStr.split(",")) + .map(String::trim) + .map(String::toLowerCase) + .map(DataType::fromString) + .filter(Objects::nonNull) + .collect(Collectors.toSet()); + } + + private Set parseFormats(String formatsStr) { + if (StringUtils.isBlank(formatsStr)) { + return EnumSet.of(SerializationFormat.AVRO); + } + return Arrays.stream(formatsStr.split(",")) + .map(String::trim) + .map(String::toLowerCase) + .map(SerializationFormat::fromString) + .filter(Objects::nonNull) + .collect(Collectors.toSet()); + } + + public long getRecordsCount() { + return recordsCount; + } + public int getFieldCountPerRecord() { + return fieldCountPerRecord; + } + public int getPayloadsCount() { + return payloadsCount; + } + public int getBatchSizeBytes() { + return batchSizeBytes; + } + public Set getEnabledDataTypes() { + return enabledDataTypes; + } + public Set getEnabledFormats() { + return enabledFormats; + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/PerfTestCase.java b/core/src/main/java/kafka/automq/table/perf/PerfTestCase.java new file mode 100644 index 0000000000..af82f583c7 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/PerfTestCase.java @@ -0,0 +1,233 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import kafka.automq.table.process.DefaultRecordProcessor; +import kafka.automq.table.process.RecordProcessor; +import kafka.automq.table.process.convert.RawConverter; +import kafka.automq.table.process.transform.FlattenTransform; +import kafka.automq.table.worker.IcebergTableManager; +import kafka.automq.table.worker.IcebergWriter; +import kafka.automq.table.worker.WorkerConfig; + +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.record.TimestampType; + +import com.google.common.collect.ImmutableMap; + +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +public abstract class PerfTestCase { + // Cache InMemoryFileIO files map to avoid repeated reflection overhead + private static final Map IN_MEMORY_FILES; + + static { + Map files; + try { + Class clazz = Class.forName("org.apache.iceberg.inmemory.InMemoryFileIO"); + java.lang.reflect.Field field = clazz.getDeclaredField("IN_MEMORY_FILES"); + field.setAccessible(true); + @SuppressWarnings("unchecked") + Map f = (Map) field.get(null); + files = f; + } catch (Exception e) { + // Fallback to empty map; clear operation becomes a no-op + files = new java.util.HashMap<>(); + } + IN_MEMORY_FILES = files; + } + + protected final DataType dataType; + protected final String formatName; + protected final int fieldCount; + protected final int payloadCount; + protected final int batchSizeBytes; + + public PerfTestCase(DataType dataType, String formatName, int fieldCount, int payloadCount, PerfConfig config) { + this.dataType = dataType; + this.formatName = formatName; + this.fieldCount = fieldCount; + this.payloadCount = payloadCount; + this.batchSizeBytes = config.getBatchSizeBytes(); + } + + abstract PayloadManager getPayloadManager(); + + public BenchmarkResult runBenchmark(long recordsCount) { + try { + // Simple warmup + runTest(100_000); + + // Actual test + long startTime = System.nanoTime(); + long fieldCount = runTest(recordsCount); + long durationNs = System.nanoTime() - startTime; + + return BenchmarkResult.success(formatName, dataType.getName(), durationNs, recordsCount, fieldCount); + } catch (Exception e) { + return BenchmarkResult.failure(formatName, dataType.getName(), e.getMessage()); + } + } + + private long runTest(long recordsCount) throws IOException { + TableIdentifier tableId = TableIdentifier.parse("test.benchmark"); + WorkerConfig workerConfig = new BenchmarkWorkerConfig(); + IcebergWriter writer = null; + int currentBatchSize = 0; + final int batchSizeLimit = this.batchSizeBytes; + long totalFieldCount = 0; + + for (long i = 0; i < recordsCount; i++) { + if (writer == null) { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("test", ImmutableMap.of()); + RecordProcessor processor = new DefaultRecordProcessor("", RawConverter.INSTANCE, createConverter(), List.of(FlattenTransform.INSTANCE)); + writer = new IcebergWriter(new IcebergTableManager(catalog, tableId, workerConfig), processor, workerConfig); + writer.setOffset(0, i); + } + + byte[] payload = getPayloadManager().nextPayload(); + currentBatchSize += payload.length; + + writer.write(0, new SimpleRecord(i, payload)); + + if (currentBatchSize > batchSizeLimit) { + totalFieldCount += finalizeWriter(writer); + writer = null; + currentBatchSize = 0; + } + } + + if (writer != null) { + totalFieldCount += finalizeWriter(writer); + } + return totalFieldCount; + } + + public static void clearInMemoryFiles() { + try { + IN_MEMORY_FILES.clear(); + } catch (Exception ignored) { + // Ignore cleanup failures + } + } + + private long finalizeWriter(IcebergWriter writer) throws IOException { + try { + writer.complete(); + return writer.topicMetric().fieldCount(); + } finally { + clearInMemoryFiles(); + } + } + + protected abstract byte[] generatePayload(); + protected abstract kafka.automq.table.process.Converter createConverter(); + + static class SimpleRecord implements org.apache.kafka.common.record.Record { + final long offset; + final byte[] value; + + public SimpleRecord(long offset, byte[] value) { + this.offset = offset; + this.value = value; + } + + @Override + public long offset() { + return offset; + } + + @Override + public int sequence() { + return 0; + } + + @Override + public int sizeInBytes() { + return value.length; + } + + @Override + public long timestamp() { + return 0; + } + + @Override + public void ensureValid() { + + } + + @Override + public int keySize() { + return 0; + } + + @Override + public boolean hasKey() { + return false; + } + + @Override + public ByteBuffer key() { + return null; + } + + @Override + public int valueSize() { + return 0; + } + + @Override + public boolean hasValue() { + return true; + } + + @Override + public ByteBuffer value() { + return ByteBuffer.wrap(value); + } + + @Override + public boolean hasMagic(byte b) { + return false; + } + + @Override + public boolean isCompressed() { + return false; + } + + @Override + public boolean hasTimestampType(TimestampType type) { + return false; + } + + @Override + public Header[] headers() { + return new Header[0]; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/ProtobufTestCase.java b/core/src/main/java/kafka/automq/table/perf/ProtobufTestCase.java new file mode 100644 index 0000000000..d8b685f39d --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/ProtobufTestCase.java @@ -0,0 +1,152 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; +import kafka.automq.table.deserializer.proto.schema.MessageDefinition; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.convert.ProtobufRegistryConverter; + +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Message; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +class ProtobufTestCase extends PerfTestCase { + private final Descriptors.Descriptor descriptor; + private final PayloadManager payloadManager; + + public ProtobufTestCase(DataType dataType, int fieldCount, int payloadCount, PerfConfig config) { + super(dataType, "proto", fieldCount, payloadCount, config); + this.descriptor = createDescriptor(dataType, fieldCount); + this.payloadManager = new PayloadManager(() -> generatePayloadWithDescriptor(descriptor), payloadCount); + } + + @Override + PayloadManager getPayloadManager() { + return payloadManager; + } + + @Override + protected byte[] generatePayload() { + return generatePayloadWithDescriptor(descriptor); + } + + private byte[] generatePayloadWithDescriptor(Descriptors.Descriptor descriptorToUse) { + DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptorToUse); + fillBuilder(builder, dataType, fieldCount); + return createProtoBufValue(builder.build()); + } + + @Override + protected Converter createConverter() { + StaticProtobufDeserializer deserializer = new StaticProtobufDeserializer(descriptor); + return new ProtobufRegistryConverter(deserializer); + } + + private Descriptors.Descriptor createDescriptor(DataType dataType, int fieldCount) { + DynamicSchema.Builder schemaBuilder = DynamicSchema.newBuilder(); + schemaBuilder.setPackage("example"); + schemaBuilder.setName("DynamicTest.proto"); + + MessageDefinition.Builder msgDefBuilder = MessageDefinition.newBuilder("TestMessage"); + + for (int i = 0; i < fieldCount; i++) { + String fieldName = "f" + i; + int fieldNumber = i + 1; + addFieldToMessage(msgDefBuilder, fieldName, fieldNumber, dataType); + } + + schemaBuilder.addMessageDefinition(msgDefBuilder.build()); + + try { + DynamicSchema schema = schemaBuilder.build(); + return schema.getMessageDescriptor("TestMessage"); + } catch (Exception e) { + throw new RuntimeException("Schema build failed", e); + } + } + + private void addFieldToMessage(MessageDefinition.Builder msgDefBuilder, String fieldName, int fieldNumber, DataType dataType) { + switch (dataType) { + case BOOLEAN -> msgDefBuilder.addField("required", "bool", fieldName, fieldNumber, null); + case INT -> msgDefBuilder.addField("required", "int32", fieldName, fieldNumber, null); + case LONG, TIMESTAMP -> msgDefBuilder.addField("required", "int64", fieldName, fieldNumber, null); + case DOUBLE -> msgDefBuilder.addField("required", "double", fieldName, fieldNumber, null); + case STRING -> msgDefBuilder.addField("required", "string", fieldName, fieldNumber, null); + case BINARY -> msgDefBuilder.addField("required", "bytes", fieldName, fieldNumber, null); + case NESTED -> { + MessageDefinition nested = MessageDefinition.newBuilder("NestedType" + fieldName) + .addField("required", "bool", "nf1", 1, null) + .build(); + msgDefBuilder.addMessageDefinition(nested); + msgDefBuilder.addField("required", "NestedType" + fieldName, fieldName, fieldNumber, null); + } + case ARRAY -> msgDefBuilder.addField("repeated", "bool", fieldName, fieldNumber, null); + } + } + + private void fillBuilder(DynamicMessage.Builder builder, DataType dataType, int fieldCount) { + for (int i = 0; i < fieldCount; i++) { + String fieldName = "f" + i; + setFieldValue(builder, fieldName, dataType); + } + } + + private void setFieldValue(DynamicMessage.Builder builder, String fieldName, DataType dataType) { + Descriptors.FieldDescriptor field = builder.getDescriptorForType().findFieldByName(fieldName); + + switch (dataType) { + case BOOLEAN, INT, LONG, DOUBLE, TIMESTAMP, STRING -> + builder.setField(field, dataType.generateValue()); + case BINARY -> { + byte[] bytes = new byte[32]; + java.util.concurrent.ThreadLocalRandom.current().nextBytes(bytes); + builder.setField(field, bytes); + } + case NESTED -> { + Descriptors.Descriptor nestedDescriptor = builder.getDescriptorForType().findNestedTypeByName("NestedType" + fieldName); + DynamicMessage nestedMessage = DynamicMessage.newBuilder(nestedDescriptor) + .setField(nestedDescriptor.findFieldByName("nf1"), DataType.BOOLEAN.generateValue()) + .build(); + builder.setField(field, nestedMessage); + } + case ARRAY -> + builder.addRepeatedField(field, DataType.BOOLEAN.generateValue()); + } + } + + private static byte[] createProtoBufValue(Message message) { + try { + byte[] protobufBytes = message.toByteArray(); + + ByteBuf buf = Unpooled.buffer(1 + 4 + protobufBytes.length); + buf.writeByte((byte) 0x0); + buf.writeInt(0); + buf.writeBytes(protobufBytes); + + return buf.array(); + } catch (Throwable e) { + throw new RuntimeException(e); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/perf/SerializationFormat.java b/core/src/main/java/kafka/automq/table/perf/SerializationFormat.java new file mode 100644 index 0000000000..f36268eb6d --- /dev/null +++ b/core/src/main/java/kafka/automq/table/perf/SerializationFormat.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.perf; + +public enum SerializationFormat { + AVRO("avro"), + PROTOBUF("proto"); + + private final String name; + + SerializationFormat(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public static SerializationFormat fromString(String name) { + for (SerializationFormat format : values()) { + if (format.name.equals(name)) { + return format; + } + } + return null; + } +} diff --git a/core/src/main/java/kafka/automq/table/process/ConversionResult.java b/core/src/main/java/kafka/automq/table/process/ConversionResult.java new file mode 100644 index 0000000000..fc891ace82 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/ConversionResult.java @@ -0,0 +1,57 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; + +import java.util.Objects; + + +public final class ConversionResult { + + private final Object value; + private final Schema schema; + private final String schemaIdentity; + + public ConversionResult(Object value, Schema schema, String schemaIdentity) { + this.value = value; + this.schema = Objects.requireNonNull(schema, "schema cannot be null"); + this.schemaIdentity = Objects.requireNonNull(schemaIdentity, "schemaIdentity cannot be null"); + } + + public ConversionResult(GenericRecord record, String schemaIdentity) { + this.value = Objects.requireNonNull(record, "record cannot be null"); + this.schema = record.getSchema(); + this.schemaIdentity = Objects.requireNonNull(schemaIdentity, "schemaIdentity cannot be null"); + } + + public Object getValue() { + return value; + } + + public Schema getSchema() { + return schema; + } + + public String getSchemaIdentity() { + return schemaIdentity; + } +} diff --git a/core/src/main/java/kafka/automq/table/process/Converter.java b/core/src/main/java/kafka/automq/table/process/Converter.java new file mode 100644 index 0000000000..672d2380e1 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/Converter.java @@ -0,0 +1,31 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.exception.ConverterException; + +import java.nio.ByteBuffer; + + +public interface Converter { + + ConversionResult convert(String topic, ByteBuffer buffer) throws ConverterException; + +} diff --git a/core/src/main/java/kafka/automq/table/process/DataError.java b/core/src/main/java/kafka/automq/table/process/DataError.java new file mode 100644 index 0000000000..ddf2b97259 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/DataError.java @@ -0,0 +1,173 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import java.util.Objects; + +/** + * Represents a data-level error that occurred during record processing. + * + *

This class encapsulates recoverable errors that occur when processing + * individual records. Unlike system-level exceptions, these errors allow + * for more granular error handling strategies, such as:

+ *
    + *
  • Skipping bad records while continuing processing
    • + *
  • Logging detailed error information for later analysis
    • + *
  • Implementing custom retry logic for specific error types
    • + *
  • Routing failed records to dead letter queues
    • + *
+ * + *

DataError instances are immutable and provide comprehensive error + * information including type classification, human-readable messages, + * and underlying causes.

+ * + * @see ProcessingResult + * @see ErrorType + */ +public final class DataError { + + private final ErrorType type; + private final String message; + private final Throwable cause; + + /** + * Creates a new DataError with the specified type, message, and underlying cause. + * + * @param type the classification of this error + * @param message a human-readable description of the error + * @param cause the underlying exception that caused this error (may be null) + * @throws IllegalArgumentException if type or message is null + */ + public DataError(ErrorType type, String message, Throwable cause) { + this.type = Objects.requireNonNull(type, "ErrorType cannot be null"); + this.message = Objects.requireNonNull(message, "Error message cannot be null"); + this.cause = cause; + } + + /** + * Returns the classification of this error. + * + * @return the error type, never null + */ + public ErrorType getType() { + return type; + } + + /** + * Returns a human-readable description of this error. + * + * @return the error message, never null + */ + public String getMessage() { + return message; + } + + /** + * Returns the underlying exception that caused this error. + * + * @return the cause exception, or null if there is no underlying cause + */ + public Throwable getCause() { + return cause; + } + + /** + * Checks if this error has an underlying cause exception. + * + * @return true if there is an underlying cause, false otherwise + */ + public boolean hasCause() { + return cause != null; + } + + /** + * Creates a detailed error message that includes type, message, and cause information. + * + * @return a comprehensive error description + */ + public String getDetailedMessage() { + StringBuilder sb = new StringBuilder(); + sb.append("[").append(type).append("] ").append(message); + if (hasCause()) { + sb.append(" (caused by: ").append(cause.getClass().getSimpleName()); + if (cause.getMessage() != null) { + sb.append(" - ").append(cause.getMessage()); + } + sb.append(")"); + } + return sb.toString(); + } + + @Override + public boolean equals(Object obj) { + if (this == obj) return true; + if (obj == null || getClass() != obj.getClass()) return false; + + DataError dataError = (DataError) obj; + return type == dataError.type && + Objects.equals(message, dataError.message) && + Objects.equals(cause, dataError.cause); + } + + @Override + public int hashCode() { + return Objects.hash(type, message, cause); + } + + @Override + public String toString() { + return "DataError{" + + "type=" + type + + ", message='" + message + '\'' + + ", cause=" + (cause != null ? cause.getClass().getSimpleName() : "null") + + '}'; + } + + /** + * Classification of data processing errors for appropriate error handling. + */ + public enum ErrorType { + /** + * Indicates that the record's data is malformed or invalid. + * This can be due to a null payload, incorrect data size, or a missing or unknown magic byte that prevents further processing. + */ + DATA_ERROR, + /** + * Indicates a failure during the data deserialization or conversion step. + * This can be caused by an inability to fetch a schema from a registry, serialization errors (e.g., Avro, Protobuf), or other schema-related failures. + */ + CONVERT_ERROR, + /** + * Indicates a failure within a data transformation step. + * This can be due to the input data not matching the format expected by a transform (e.g., an invalid Debezium record). + */ + TRANSFORMATION_ERROR, + + /** + * + */ + SYSTEM_ERROR, + + /** + * Unknown exception during processing. + */ + UNKNOW_ERROR + } +} diff --git a/core/src/main/java/kafka/automq/table/process/DefaultRecordProcessor.java b/core/src/main/java/kafka/automq/table/process/DefaultRecordProcessor.java new file mode 100644 index 0000000000..a1d7ce8d96 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/DefaultRecordProcessor.java @@ -0,0 +1,252 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.exception.ConverterException; +import kafka.automq.table.process.exception.InvalidDataException; +import kafka.automq.table.process.exception.RecordProcessorException; +import kafka.automq.table.process.exception.SchemaRegistrySystemException; +import kafka.automq.table.process.exception.TransformException; + +import org.apache.kafka.common.cache.Cache; +import org.apache.kafka.common.cache.LRUCache; +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.record.Record; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.jetbrains.annotations.NotNull; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; + +import static java.net.HttpURLConnection.HTTP_NOT_FOUND; +import static kafka.automq.table.process.RecordAssembler.KAFKA_VALUE_FIELD; +import static kafka.automq.table.process.RecordAssembler.ensureOptional; + +/** + * Default implementation of RecordProcessor using a two-stage processing pipeline. + * + * @see RecordProcessor + * @see Converter + * @see Transform + */ +public class DefaultRecordProcessor implements RecordProcessor { + private static final Schema HEADER_SCHEMA = Schema.createMap(Schema.create(Schema.Type.BYTES)); + private static final String HEADER_SCHEMA_IDENTITY = String.valueOf(HEADER_SCHEMA.hashCode()); + private static final ConversionResult EMPTY_HEADERS_RESULT = + new ConversionResult(Map.of(), HEADER_SCHEMA, HEADER_SCHEMA_IDENTITY); + private final String topicName; + private final Converter keyConverter; + private final Converter valueConverter; + private final List transformChain; + private final RecordAssembler recordAssembler; // Reusable assembler + private final String transformIdentity; // precomputed transform chain identity + + private static final int VALUE_WRAPPER_SCHEMA_CACHE_MAX = 32; + private final Cache valueWrapperSchemaCache = new LRUCache<>(VALUE_WRAPPER_SCHEMA_CACHE_MAX); + + public DefaultRecordProcessor(String topicName, Converter keyConverter, Converter valueConverter) { + this.transformChain = new ArrayList<>(); + this.topicName = topicName; + this.keyConverter = keyConverter; + this.valueConverter = valueConverter; + this.recordAssembler = new RecordAssembler(); + this.transformIdentity = ""; // no transforms + } + + public DefaultRecordProcessor(String topicName, Converter keyConverter, Converter valueConverter, List transforms) { + this.transformChain = transforms; + this.topicName = topicName; + this.keyConverter = keyConverter; + this.valueConverter = valueConverter; + this.recordAssembler = new RecordAssembler(); + + // Precompute transform identity (names joined by comma) + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < this.transformChain.size(); i++) { + if (i > 0) sb.append(','); + sb.append(this.transformChain.get(i).getName()); + } + this.transformIdentity = sb.toString(); + } + + @Override + public ProcessingResult process(int partition, Record kafkaRecord) { + try { + Objects.requireNonNull(kafkaRecord, "Kafka record cannot be null"); + + ConversionResult headerResult = processHeaders(kafkaRecord); + ConversionResult keyResult = keyConverter.convert(topicName, kafkaRecord.key()); + ConversionResult valueResult = valueConverter.convert(topicName, kafkaRecord.value()); + + GenericRecord baseRecord = wrapValue(valueResult); + GenericRecord transformedRecord = applyTransformChain(baseRecord, partition, kafkaRecord); + + String schemaIdentity = generateCompositeSchemaIdentity(headerResult, keyResult, valueResult); + + GenericRecord record = recordAssembler + .reset(transformedRecord) + .withHeader(headerResult) + .withKey(keyResult) + .withSchemaIdentity(schemaIdentity) + .withMetadata(partition, kafkaRecord.offset(), kafkaRecord.timestamp()) + .assemble(); + Schema schema = record.getSchema(); + + return new ProcessingResult(record, schema, schemaIdentity); + } catch (ConverterException e) { + return getProcessingResult(kafkaRecord, "Convert operation failed for record: %s", DataError.ErrorType.CONVERT_ERROR, e); + } catch (TransformException e) { + return getProcessingResult(kafkaRecord, "Transform operation failed for record: %s", DataError.ErrorType.TRANSFORMATION_ERROR, e); + } catch (InvalidDataException e) { + return getProcessingResult(kafkaRecord, "Transform operation failed for record: %s", DataError.ErrorType.DATA_ERROR, e); + } catch (Exception e) { + if (e.getCause() instanceof RestClientException) { + RestClientException exception = (RestClientException) e.getCause(); + // io.confluent.kafka.serializers.AbstractKafkaSchemaSerDe#toKafkaException + if (isSchemaOrSubjectNotFoundException(exception)) { // not found + return getProcessingResult(kafkaRecord, "Schema or subject not found for record: %s", DataError.ErrorType.CONVERT_ERROR, exception); + } + throw SchemaRegistrySystemException.fromStatusCode(exception, buildRecordContext(kafkaRecord)); + } + return getProcessingResult(kafkaRecord, "Unexpected error processing record: %s", DataError.ErrorType.UNKNOW_ERROR, e); + } + } + + // io.confluent.kafka.schemaregistry.client.CachedSchemaRegistryClient#isSchemaOrSubjectNotFoundException + private boolean isSchemaOrSubjectNotFoundException(RestClientException rce) { + return rce.getStatus() == HTTP_NOT_FOUND + && (rce.getErrorCode() == 40403 // SCHEMA_NOT_FOUND_ERROR_CODE + || rce.getErrorCode() == 40401); // SUBJECT_NOT_FOUND_ERROR_CODE + } + + + @NotNull + private ProcessingResult getProcessingResult(Record kafkaRecord, String format, DataError.ErrorType unknow, Exception e) { + String recordContext = buildRecordContext(kafkaRecord); + String errorMsg = String.format(format, recordContext); + DataError error = new DataError(unknow, errorMsg + ": " + e.getMessage(), e); + return new ProcessingResult(error); + } + + private ConversionResult processHeaders(Record kafkaRecord) throws ConverterException { + try { + Header[] recordHeaders = kafkaRecord.headers(); + if (recordHeaders == null || recordHeaders.length == 0) { + return EMPTY_HEADERS_RESULT; + } + + int n = recordHeaders.length; + + // Small maps: use Map.of for zero/one header handled above; for one here (defensive), use Map.of + if (n == 1) { + Header h = recordHeaders[0]; + ByteBuffer value = h.value() != null ? ByteBuffer.wrap(h.value()) : null; + Map headers = Map.of(h.key(), value); + return new ConversionResult(headers, HEADER_SCHEMA, HEADER_SCHEMA_IDENTITY); + } + + // Larger maps: pre-size HashMap + Map headers = new HashMap<>(Math.max(16, (int) (n / 0.75f) + 1)); + for (Header header : recordHeaders) { + ByteBuffer value = header.value() != null ? ByteBuffer.wrap(header.value()) : null; + headers.put(header.key(), value); + } + return new ConversionResult(headers, HEADER_SCHEMA, HEADER_SCHEMA_IDENTITY); + } catch (Exception e) { + throw new ConverterException("Failed to process headers", e); + } + } + + private GenericRecord wrapValue(ConversionResult valueResult) { + Object valueContent = valueResult.getValue(); + Schema recordSchema = valueWrapperSchemaCache.get(valueResult.getSchemaIdentity()); + if (recordSchema == null) { + Schema.Field valueField = new Schema.Field( + KAFKA_VALUE_FIELD, + ensureOptional(valueResult.getSchema()), + null, null); + Schema schema = Schema.createRecord("KafkaValueWrapper", null, "kafka.automq.table.process", false); + schema.setFields(Collections.singletonList(valueField)); + valueWrapperSchemaCache.put(valueResult.getSchemaIdentity(), schema); + recordSchema = schema; + } + + GenericRecord baseRecord = new GenericData.Record(recordSchema); + baseRecord.put(KAFKA_VALUE_FIELD, valueContent); + return baseRecord; + } + + private GenericRecord applyTransformChain(GenericRecord baseRecord, int partition, Record kafkaRecord) throws TransformException { + if (transformChain.isEmpty()) { + return baseRecord; + } + + GenericRecord currentRecord = baseRecord; + TransformContext context = new TransformContext(kafkaRecord, topicName, partition); + + for (Transform transform : transformChain) { + currentRecord = transform.apply(currentRecord, context); + if (currentRecord == null) { + throw new TransformException("Transform " + transform.getName() + " returned null record"); + } + } + return currentRecord; + } + + + private String generateCompositeSchemaIdentity( + ConversionResult headerResult, + ConversionResult keyResult, + ConversionResult valueResult) { + // Extract schema identities + String headerIdentity = headerResult.getSchemaIdentity(); + String keyIdentity = keyResult.getSchemaIdentity(); + String valueIdentity = valueResult.getSchemaIdentity(); + return "h:" + headerIdentity + "|v:" + valueIdentity + "|k:" + keyIdentity + "|t:" + transformIdentity; + } + + @Override + public void configure(Map configs) throws RecordProcessorException { + // ignore + } + + private String buildRecordContext(org.apache.kafka.common.record.Record kafkaRecord) { + return String.format("topic=%s, key=%s, offset=%d, timestamp=%d", + topicName, + kafkaRecord.key(), + kafkaRecord.offset(), + kafkaRecord.timestamp()); + } + + public String getTopicName() { + return topicName; + } + +} diff --git a/core/src/main/java/kafka/automq/table/process/ProcessingResult.java b/core/src/main/java/kafka/automq/table/process/ProcessingResult.java new file mode 100644 index 0000000000..4ef5882530 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/ProcessingResult.java @@ -0,0 +1,116 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; + +import java.util.Objects; + +/** + * Comprehensive result object returned by RecordProcessor operations. + * + *
    + *
  • Success: Contains finalRecord, finalSchema, and + * finalSchemaIdentity; error is null
    • + *
  • Failure: Contains detailed error information; + * all other fields are null
    • + *
+ * + * @see RecordProcessor#process(int, org.apache.kafka.common.record.Record) + * @see DataError + */ +public final class ProcessingResult { + + private final GenericRecord finalRecord; + private final Schema finalSchema; + private final String finalSchemaIdentity; + private final DataError error; + + /** + * Creates a successful processing result. + * + * @param finalRecord the Avro GenericRecord ready for further processing, must not be null + * @param finalSchema the Avro schema matching the finalRecord, must not be null + * @param finalSchemaIdentity unique identifier for schema comparison, must not be null + * @throws IllegalArgumentException if any parameter is null + */ + public ProcessingResult(GenericRecord finalRecord, Schema finalSchema, String finalSchemaIdentity) { + this.finalRecord = Objects.requireNonNull(finalRecord, "finalRecord cannot be null"); + this.finalSchema = Objects.requireNonNull(finalSchema, "finalSchema cannot be null"); + this.finalSchemaIdentity = Objects.requireNonNull(finalSchemaIdentity, "finalSchemaIdentity cannot be null"); + this.error = null; + } + + /** + * Creates a failed processing result. + * + * @param error the data error that occurred during processing, must not be null + * @throws IllegalArgumentException if error is null + */ + public ProcessingResult(DataError error) { + this.finalRecord = null; + this.finalSchema = null; + this.finalSchemaIdentity = null; + this.error = Objects.requireNonNull(error, "error cannot be null"); + } + + public GenericRecord getFinalRecord() { + return finalRecord; + } + public Schema getFinalSchema() { + return finalSchema; + } + public String getFinalSchemaIdentity() { + return finalSchemaIdentity; + } + public DataError getError() { + return error; + } + public boolean isSuccess() { + return error == null; + } + + @Override + public boolean equals(Object obj) { + if (this == obj) return true; + if (obj == null || getClass() != obj.getClass()) return false; + + ProcessingResult that = (ProcessingResult) obj; + return Objects.equals(finalRecord, that.finalRecord) && + Objects.equals(finalSchema, that.finalSchema) && + Objects.equals(finalSchemaIdentity, that.finalSchemaIdentity) && + Objects.equals(error, that.error); + } + + @Override + public int hashCode() { + return Objects.hash(finalRecord, finalSchema, finalSchemaIdentity, error); + } + + @Override + public String toString() { + if (!isSuccess()) { + return "ProcessingResult{success=true, schemaIdentity=" + finalSchemaIdentity + "}"; + } else { + return "ProcessingResult{success=false, error=" + error + "}"; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/RecordAssembler.java b/core/src/main/java/kafka/automq/table/process/RecordAssembler.java new file mode 100644 index 0000000000..64750b5a7f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/RecordAssembler.java @@ -0,0 +1,305 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import org.apache.kafka.common.cache.LRUCache; + +import org.apache.avro.JsonProperties; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.SchemaNormalization; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.util.internal.Accessor; + +import java.util.ArrayList; +import java.util.List; + +/** + * A specialized assembler for constructing the final record structure + * in a clean, fluent manner following the builder pattern. + *

+ * This class also serves as the holder for the public contract of field names. + */ +public final class RecordAssembler { + private static final Schema NULL_SCHEMA = Schema.create(Schema.Type.NULL); + + public static final String KAFKA_HEADER_FIELD = "_kafka_header"; + public static final String KAFKA_KEY_FIELD = "_kafka_key"; + public static final String KAFKA_VALUE_FIELD = "_kafka_value"; + public static final String KAFKA_METADATA_FIELD = "_kafka_metadata"; + public static final String METADATA_PARTITION_FIELD = "partition"; + public static final String METADATA_OFFSET_FIELD = "offset"; + public static final String METADATA_TIMESTAMP_FIELD = "timestamp"; + + private static final Schema METADATA_SCHEMA = SchemaBuilder + .record("KafkaMetadata") + .namespace("kafka.automq.table.process") + .doc("Holds metadata about the original Kafka record.") + .fields() + .name(METADATA_PARTITION_FIELD).doc("Partition id").type().intType().noDefault() + .name(METADATA_OFFSET_FIELD).doc("Record offset").type().longType().noDefault() + .name(METADATA_TIMESTAMP_FIELD).doc("Record timestamp").type().longType().noDefault() + .endRecord(); + + private static final int SCHEMA_CACHE_MAX = 32; + // Cache of assembled schema + precomputed indexes bound to a schema identity + private final LRUCache assemblerSchemaCache = new LRUCache<>(SCHEMA_CACHE_MAX); + + // Reusable state - reset for each record + private GenericRecord baseRecord; + private ConversionResult headerResult; + private ConversionResult keyResult; + private int partition; + private long offset; + private long timestamp; + private String schemaIdentity; + + public RecordAssembler() { + } + + public RecordAssembler reset(GenericRecord baseRecord) { + this.baseRecord = baseRecord; + this.headerResult = null; + this.keyResult = null; + this.partition = 0; + this.offset = 0L; + this.timestamp = 0L; + this.schemaIdentity = null; + return this; + } + + public RecordAssembler withHeader(ConversionResult headerResult) { + this.headerResult = headerResult; + return this; + } + + public RecordAssembler withKey(ConversionResult keyResult) { + this.keyResult = keyResult; + return this; + } + + + public RecordAssembler withSchemaIdentity(String schemaIdentity) { + this.schemaIdentity = schemaIdentity; + return this; + } + + public RecordAssembler withMetadata(int partition, long offset, long timestamp) { + this.partition = partition; + this.offset = offset; + this.timestamp = timestamp; + return this; + } + + public GenericRecord assemble() { + AssemblerSchema aSchema = getOrCreateAssemblerSchema(); + // Return a lightweight view that implements GenericRecord + // and adapts schema position/name lookups to the underlying values + // without copying the base record data. + return new AssembledRecordView(aSchema, baseRecord, + headerResult != null ? headerResult.getValue() : null, + keyResult != null ? keyResult.getValue() : null, + partition, offset, timestamp); + } + + private AssemblerSchema getOrCreateAssemblerSchema() { + if (schemaIdentity == null) { + long baseFp = SchemaNormalization.parsingFingerprint64(baseRecord.getSchema()); + long keyFp = keyResult != null ? SchemaNormalization.parsingFingerprint64(keyResult.getSchema()) : 0L; + long headerFp = headerResult != null ? SchemaNormalization.parsingFingerprint64(headerResult.getSchema()) : 0L; + long metadataFp = SchemaNormalization.parsingFingerprint64(METADATA_SCHEMA); + + schemaIdentity = "v:" + Long.toUnsignedString(baseFp) + + "|k:" + Long.toUnsignedString(keyFp) + + "|h:" + Long.toUnsignedString(headerFp) + + "|m:" + Long.toUnsignedString(metadataFp); + } + final String cacheKey = schemaIdentity; + AssemblerSchema cached = assemblerSchemaCache.get(cacheKey); + if (cached != null) { + return cached; + } + AssemblerSchema created = buildFinalAssemblerSchema(); + assemblerSchemaCache.put(cacheKey, created); + return created; + } + + private AssemblerSchema buildFinalAssemblerSchema() { + List finalFields = new ArrayList<>(baseRecord.getSchema().getFields().size() + 3); + Schema baseSchema = baseRecord.getSchema(); + for (Schema.Field field : baseSchema.getFields()) { + // Accessor keeps the original Schema instance (preserving logical types) while skipping default-value revalidation. + Schema.Field f = Accessor.createField(field.name(), field.schema(), field.doc(), Accessor.defaultValue(field), false, field.order()); + finalFields.add(f); + } + + int baseFieldCount = baseSchema.getFields().size(); + int headerIndex = -1; + int keyIndex = -1; + int metadataIndex = -1; + + if (headerResult != null) { + Schema optionalHeaderSchema = ensureOptional(headerResult.getSchema()); + finalFields.add(new Schema.Field(KAFKA_HEADER_FIELD, optionalHeaderSchema, "Kafka record headers", JsonProperties.NULL_VALUE)); + headerIndex = baseFieldCount; + } + if (keyResult != null) { + Schema optionalKeySchema = ensureOptional(keyResult.getSchema()); + finalFields.add(new Schema.Field(KAFKA_KEY_FIELD, optionalKeySchema, "Kafka record key", JsonProperties.NULL_VALUE)); + keyIndex = (headerIndex >= 0) ? baseFieldCount + 1 : baseFieldCount; + } + + Schema optionalMetadataSchema = ensureOptional(METADATA_SCHEMA); + finalFields.add(new Schema.Field(KAFKA_METADATA_FIELD, optionalMetadataSchema, "Kafka record metadata", JsonProperties.NULL_VALUE)); + metadataIndex = baseFieldCount + (headerIndex >= 0 ? 1 : 0) + (keyIndex >= 0 ? 1 : 0); + + Schema finalSchema = Schema.createRecord(baseSchema.getName() + "WithMetadata", null, + "kafka.automq.table.process", false, finalFields); + + return new AssemblerSchema(finalSchema, baseFieldCount, headerIndex, keyIndex, metadataIndex); + } + + public static Schema ensureOptional(Schema schema) { + if (schema.getType() == Schema.Type.UNION) { + boolean hasNull = false; + List types = schema.getTypes(); + for (Schema type : types) { + if (type.getType() == Schema.Type.NULL) { + hasNull = true; + break; + } + } + if (hasNull) { + return schema; + } + List withNull = new ArrayList<>(types.size() + 1); + withNull.add(NULL_SCHEMA); + withNull.addAll(types); + return Schema.createUnion(withNull); + } + return Schema.createUnion(List.of(NULL_SCHEMA, schema)); + } + + /** + * A read-only GenericRecord view that adapts accesses (by name or position) + * to the underlying base record and the synthetic kafka fields. + */ + private static final class AssembledRecordView implements GenericRecord { + private final Schema finalSchema; + private final int finalFieldCount; + private final GenericRecord baseRecord; + private final Object headerValue; // May be null if not present in schema + private final Object keyValue; // May be null if not present in schema + private final int baseFieldCount; + private final int headerIndex; // -1 if absent + private final int keyIndex; // -1 if absent + private final int metadataIndex; // always >= 0 + + private GenericRecord metadataRecord; + + AssembledRecordView(AssemblerSchema aSchema, + GenericRecord baseRecord, + Object headerValue, + Object keyValue, + int partition, + long offset, + long timestamp) { + this.finalSchema = aSchema.schema; + this.finalFieldCount = finalSchema.getFields().size(); + this.baseRecord = baseRecord; + this.headerValue = headerValue; + this.keyValue = keyValue; + + this.baseFieldCount = aSchema.baseFieldCount; + this.headerIndex = aSchema.headerIndex; + this.keyIndex = aSchema.keyIndex; + this.metadataIndex = aSchema.metadataIndex; + + this.metadataRecord = new GenericData.Record(METADATA_SCHEMA); + metadataRecord.put(METADATA_PARTITION_FIELD, partition); + metadataRecord.put(METADATA_OFFSET_FIELD, offset); + metadataRecord.put(METADATA_TIMESTAMP_FIELD, timestamp); + } + + @Override + public void put(String key, Object v) { + throw new UnsupportedOperationException("AssembledRecordView is read-only"); + } + + @Override + public Object get(String key) { + Schema.Field field = finalSchema.getField(key); + if (field == null) { + return null; + } + return get(field.pos()); + } + + @Override + public Schema getSchema() { + return finalSchema; + } + + @Override + public void put(int i, Object v) { + throw new UnsupportedOperationException("AssembledRecordView is read-only"); + } + + @Override + public Object get(int i) { + if (i < 0 || i >= finalFieldCount) { + throw new IndexOutOfBoundsException("Field position out of bounds: " + i); + } + // Base fields delegate directly + if (i < baseFieldCount) { + return baseRecord.get(i); + } + // Synthetic fields + if (i == headerIndex) { + return headerValue; + } + if (i == keyIndex) { + return keyValue; + } + if (i == metadataIndex) { + return metadataRecord; + } + // Should not happen if schema is consistent + return null; + } + } + + private static final class AssemblerSchema { + final Schema schema; + final int baseFieldCount; + final int headerIndex; + final int keyIndex; + final int metadataIndex; + + AssemblerSchema(Schema schema, int baseFieldCount, int headerIndex, int keyIndex, int metadataIndex) { + this.schema = schema; + this.baseFieldCount = baseFieldCount; + this.headerIndex = headerIndex; + this.keyIndex = keyIndex; + this.metadataIndex = metadataIndex; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/RecordProcessor.java b/core/src/main/java/kafka/automq/table/process/RecordProcessor.java new file mode 100644 index 0000000000..c910667673 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/RecordProcessor.java @@ -0,0 +1,45 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.exception.RecordProcessorException; + +import org.apache.kafka.common.record.Record; + +import java.util.Map; + +/** + * Processes Kafka records into standardized Avro format. + * Handles format conversion, transformations, and error processing. + */ +public interface RecordProcessor { + + /** + * Processes a Kafka record into Avro format. + * + * @param partition topic partition + * @param record the Kafka record to process + * @return ProcessingResult containing the converted record or error information + */ + ProcessingResult process(int partition, Record record); + + + void configure(Map configs) throws RecordProcessorException; +} diff --git a/core/src/main/java/kafka/automq/table/process/RecordProcessorFactory.java b/core/src/main/java/kafka/automq/table/process/RecordProcessorFactory.java new file mode 100644 index 0000000000..9b5a5106fd --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/RecordProcessorFactory.java @@ -0,0 +1,79 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.convert.ConverterFactory; +import kafka.automq.table.process.convert.StringConverter; +import kafka.automq.table.process.transform.DebeziumUnwrapTransform; +import kafka.automq.table.process.transform.FlattenTransform; +import kafka.automq.table.process.transform.SchemalessTransform; +import kafka.automq.table.worker.WorkerConfig; + +import org.apache.kafka.server.record.TableTopicSchemaType; +import org.apache.kafka.server.record.TableTopicTransformType; + +import java.util.List; + +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; + +public class RecordProcessorFactory { + private final ConverterFactory converterFactory; + + public RecordProcessorFactory(String registryUrl) { + this.converterFactory = new ConverterFactory(registryUrl); + } + + public RecordProcessorFactory(String registryUrl, SchemaRegistryClient client) { + this.converterFactory = new ConverterFactory(registryUrl, client); + } + + public RecordProcessor create(WorkerConfig config, String topic) { + // Handle deprecated configurations + if (config.schemaType() == TableTopicSchemaType.SCHEMALESS) { + return new DefaultRecordProcessor(topic, StringConverter.INSTANCE, StringConverter.INSTANCE, List.of(new SchemalessTransform())); + } + if (config.schemaType() == TableTopicSchemaType.SCHEMA) { + return new DefaultRecordProcessor(topic, + StringConverter.INSTANCE, + converterFactory.createForSchemaId(topic, false), List.of(FlattenTransform.INSTANCE)); + } + + var keyConverter = converterFactory.createKeyConverter(topic, config); + var valueConverter = converterFactory.createValueConverter(topic, config); + + var transforms = createTransforms(config.transformType()); + + return new DefaultRecordProcessor(topic, keyConverter, valueConverter, transforms); + } + + private List createTransforms(TableTopicTransformType transformType) { + if (transformType == null || TableTopicTransformType.NONE.equals(transformType)) { + return List.of(); + } + switch (transformType) { + case FLATTEN: + return List.of(FlattenTransform.INSTANCE); + case FLATTEN_DEBEZIUM: + return List.of(FlattenTransform.INSTANCE, DebeziumUnwrapTransform.INSTANCE); + default: + throw new IllegalArgumentException("Unsupported transform type: " + transformType); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/SchemaFormat.java b/core/src/main/java/kafka/automq/table/process/SchemaFormat.java new file mode 100644 index 0000000000..d61fc38ba1 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/SchemaFormat.java @@ -0,0 +1,43 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +public enum SchemaFormat { + AVRO, + JSON, + PROTOBUF; + + public static SchemaFormat fromString(String format) { + switch (format) { + case "AVRO": { + return AVRO; + } + case "JSON": { + return JSON; + } + case "PROTOBUF": { + return PROTOBUF; + } + default: { + throw new IllegalArgumentException("Unsupported schema format: " + format); + } + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/Transform.java b/core/src/main/java/kafka/automq/table/process/Transform.java new file mode 100644 index 0000000000..6bec2d726a --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/Transform.java @@ -0,0 +1,66 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.exception.TransformException; + +import org.apache.avro.generic.GenericRecord; + +import java.util.Map; + +/** + * Transform interface for Avro GenericRecord data processing. + * + *

Stateless, chainable transformations for content manipulation like CDC unwrapping, + * field mapping, and data enrichment. Implementations must be thread-safe.

+ * + * @see Converter + * @see TransformException + */ +public interface Transform { + + /** + * Configures the transform with operation-specific settings. + * + * @param configs configuration parameters, must not be null + * @throws IllegalArgumentException if configs is null or contains invalid values + */ + void configure(Map configs); + + /** + * Applies transformation to the input GenericRecord. + * + * @param record the input GenericRecord, must not be null + * @param context contextual information for the transformation + * @return the transformed GenericRecord, must not be null + * @throws TransformException if transformation fails + * @throws IllegalArgumentException if record or context is null + */ + GenericRecord apply(GenericRecord record, TransformContext context) throws TransformException; + + /** + * Returns a descriptive name for this transform. + * + * @return transform name, never null + */ + default String getName() { + return this.getClass().getSimpleName(); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/TransformContext.java b/core/src/main/java/kafka/automq/table/process/TransformContext.java new file mode 100644 index 0000000000..02f24bceb9 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/TransformContext.java @@ -0,0 +1,47 @@ +package kafka.automq.table.process; + +import org.apache.kafka.common.record.Record; + +import java.util.Objects; + +/** + * Provides contextual information to {@link Transform} operations. + * + *

This immutable class acts as a container for metadata related to the record + * being processed, but separate from the record's own data. This allows transforms + * to access information like the original Kafka record's headers or the topic name + * without polluting the data record itself.

+ * + * @see Transform#apply(org.apache.avro.generic.GenericRecord, TransformContext) + */ +public final class TransformContext { + + private final Record kafkaRecord; + private final String topicName; + private final int partition; + + /** + * Constructs a new TransformContext. + * + * @param kafkaRecord the original Kafka Record, can be null if not available. + * @param topicName the name of the topic from which the record was consumed, must not be null. + * @param partition the partition number. + */ + public TransformContext(Record kafkaRecord, String topicName, int partition) { + this.kafkaRecord = kafkaRecord; + this.topicName = Objects.requireNonNull(topicName, "topicName cannot be null"); + this.partition = partition; + } + + public Record getKafkaRecord() { + return kafkaRecord; + } + + public String getTopicName() { + return topicName; + } + + public int getPartition() { + return partition; + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/AvroRegistryConverter.java b/core/src/main/java/kafka/automq/table/process/convert/AvroRegistryConverter.java new file mode 100644 index 0000000000..5e046925f5 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/AvroRegistryConverter.java @@ -0,0 +1,110 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.convert; + +import kafka.automq.table.process.ConversionResult; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.exception.ConverterException; +import kafka.automq.table.process.exception.InvalidDataException; + +import org.apache.kafka.common.serialization.Deserializer; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Map; + +import io.confluent.kafka.schemaregistry.ParsedSchema; +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; +import io.confluent.kafka.serializers.KafkaAvroDeserializer; + +/** + * Converter for Avro format data with Confluent Schema Registry support. + * + *

This converter handles Kafka records that contain Avro-serialized data + * with schema information stored in Confluent Schema Registry. It deserializes + * the binary Avro data and resolves schemas using the registry, producing + * standardized Avro GenericRecord objects for the processing pipeline.

+ * + * @see KafkaAvroDeserializer + */ +public class AvroRegistryConverter implements Converter { + private static final int SCHEMA_ID_SIZE = 4; + private static final int HEADER_SIZE = SCHEMA_ID_SIZE + 1; // magic byte + schema id + private static final byte MAGIC_BYTE = 0x0; + + private final SchemaRegistryClient client; + private final Deserializer deserializer; + + public AvroRegistryConverter(SchemaRegistryClient client, String registryUrl, boolean isKey) { + // Initialize deserializer with the provided client + this.deserializer = new KafkaAvroDeserializer(client); + this.client = client; + // Configure the deserializer immediately upon creation + Map configs = Map.of("schema.registry.url", registryUrl); + deserializer.configure(configs, isKey); + } + + + public AvroRegistryConverter(Deserializer deserializer, SchemaRegistryClient client) { + this.deserializer = deserializer; + this.client = client; + } + + protected int getSchemaId(ByteBuffer buffer) { + if (buffer.remaining() < HEADER_SIZE) { + throw new InvalidDataException("Invalid payload size: " + buffer.remaining() + ", expected at least " + HEADER_SIZE); + } + ByteBuffer buf = buffer.duplicate(); + byte magicByte = buf.get(); + if (magicByte != MAGIC_BYTE) { + throw new InvalidDataException("Unknown magic byte: " + magicByte); + } + return buf.getInt(); + } + + @Override + public ConversionResult convert(String topic, ByteBuffer buffer) throws ConverterException { + if (buffer == null) { + throw new InvalidDataException("AvroRegistryConverter does not support null data - schema information is required"); + } + + if (buffer.remaining() == 0) { + throw new InvalidDataException("Invalid empty Avro data - schema information is required"); + } + + Object object = deserializer.deserialize(topic, null, buffer); + Schema schema; + if (object instanceof GenericRecord) { + return new ConversionResult((GenericRecord) object, String.valueOf(getSchemaId(buffer))); + } else { + try { + ParsedSchema schemaById = client.getSchemaById(getSchemaId(buffer)); + schema = (Schema) schemaById.rawSchema(); + return new ConversionResult(object, schema, String.valueOf(getSchemaId(buffer))); + } catch (RestClientException | IOException e) { + throw new ConverterException("Failed to retrieve schema from registry for topic: " + topic, e); + } + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/ConverterFactory.java b/core/src/main/java/kafka/automq/table/process/convert/ConverterFactory.java new file mode 100644 index 0000000000..ce69df497b --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/ConverterFactory.java @@ -0,0 +1,208 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.convert; + +import kafka.automq.table.deserializer.proto.LatestSchemaResolutionResolver; +import kafka.automq.table.deserializer.proto.ProtobufSchemaProvider; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.SchemaFormat; +import kafka.automq.table.process.exception.ProcessorInitializationException; +import kafka.automq.table.worker.WorkerConfig; + +import org.apache.kafka.server.record.TableTopicConvertType; + +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import io.confluent.kafka.schemaregistry.avro.AvroSchemaProvider; +import io.confluent.kafka.schemaregistry.client.CachedSchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; +import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig; + +public class ConverterFactory { + private static final Logger LOGGER = LoggerFactory.getLogger(ConverterFactory.class); + + private static final String VALUE_SUFFIX = "-value"; + private static final String KEY_SUFFIX = "-key"; + private static final String PROTOBUF_TYPE = "PROTOBUF"; + private static final Duration CACHE_EXPIRE_DURATION = Duration.ofMinutes(20); + private static final int MAX_CACHE_SIZE = 10000; + + private final String schemaRegistryUrl; + private final SchemaRegistryClient client; + private final Map converterCache = new ConcurrentHashMap<>(); + private final Cache topicSchemaFormatCache = CacheBuilder.newBuilder() + .expireAfterAccess(CACHE_EXPIRE_DURATION) + .maximumSize(MAX_CACHE_SIZE) + .build(); + + public ConverterFactory(String registryUrl) { + this.schemaRegistryUrl = registryUrl; + if (registryUrl != null && !registryUrl.trim().isEmpty()) { + this.client = new CachedSchemaRegistryClient( + registryUrl, + AbstractKafkaSchemaSerDeConfig.MAX_SCHEMAS_PER_SUBJECT_DEFAULT, + List.of(new AvroSchemaProvider(), new ProtobufSchemaProvider()), + null + ); + } else { + this.client = null; + } + } + + public ConverterFactory(String registryUrl, SchemaRegistryClient client) { + this.schemaRegistryUrl = registryUrl; + this.client = client; + } + + public Converter createKeyConverter(String topic, WorkerConfig config) { + if (topic == null || topic.trim().isEmpty()) { + throw new IllegalArgumentException("Topic cannot be null or empty"); + } + if (config == null) { + throw new IllegalArgumentException("WorkerConfig cannot be null"); + } + TableTopicConvertType convertType = config.keyConvertType(); + String subject = config.keySubject(); + String messageName = config.keyMessageFullName(); + return createConverterByType(topic, convertType, subject, messageName, true); + } + + public Converter createValueConverter(String topic, WorkerConfig config) { + if (topic == null || topic.trim().isEmpty()) { + throw new IllegalArgumentException("Topic cannot be null or empty"); + } + if (config == null) { + throw new IllegalArgumentException("WorkerConfig cannot be null"); + } + TableTopicConvertType convertType = config.valueConvertType(); + String subject = config.valueSubject(); + String messageName = config.valueMessageFullName(); + return createConverterByType(topic, convertType, subject, messageName, false); + } + + private Converter createConverterByType(String topic, TableTopicConvertType convertType, String subjectName, String messageName, boolean isKey) { + switch (convertType) { + case RAW: + return new RawConverter(); + case STRING: + return new StringConverter(); + case BY_SCHEMA_ID: + return createForSchemaId(topic, isKey); + case BY_LATEST_SCHEMA: + return createForSubjectName(topic, subjectName, messageName, isKey); + default: + throw new IllegalArgumentException("Unsupported convert type: " + convertType); + } + } + + public Converter createForSchemaId(String topic, boolean isKey) { + if (client == null) { + throw new ProcessorInitializationException("Schema Registry client is not initialized"); + } + if (topic == null || topic.trim().isEmpty()) { + throw new IllegalArgumentException("Topic cannot be null or empty"); + } + + return new LazyConverter(() -> { + String subject = getSubjectName(topic, isKey); + String schemaType = getSchemaType(subject); + SchemaFormat format = SchemaFormat.fromString(schemaType); + return converterCache.computeIfAbsent(format.name(), format1 -> createConverterForFormat(format1, isKey)); + }); + } + + public Converter createForSubjectName(String topic, String subjectName, String messageFullName, boolean isKey) { + String subject = subjectName != null ? subjectName : getSubjectName(topic); + return new LazyConverter(() -> { + String schemaType = getSchemaType(subject); + if (!PROTOBUF_TYPE.equals(schemaType)) { + throw new ProcessorInitializationException( + String.format("by_subject_name is only supported for PROTOBUF, but got %s for subject %s", schemaType, subject)); + } + + String cacheKey = schemaType + "-" + subject + "-" + (messageFullName == null ? "" : messageFullName); + return converterCache.computeIfAbsent(cacheKey, key -> { + var resolver = new LatestSchemaResolutionResolver(client, subject, messageFullName); + return new ProtobufRegistryConverter(client, schemaRegistryUrl, resolver, isKey); + }); + }); + } + + private String getSchemaType(String subject) { + if (client == null) { + throw new ProcessorInitializationException("Schema Registry client is not available"); + } + + String schemaType = topicSchemaFormatCache.getIfPresent(subject); + if (schemaType == null) { + try { + var metadata = client.getLatestSchemaMetadata(subject); + if (metadata == null) { + throw new ProcessorInitializationException("No schema found for subject: " + subject); + } + schemaType = metadata.getSchemaType(); + if (schemaType != null) { + topicSchemaFormatCache.put(subject, schemaType); + } + } catch (IOException e) { + LOGGER.error("IO error while fetching schema metadata for subject '{}'", subject, e); + throw new ProcessorInitializationException("Failed to fetch schema metadata for subject: " + subject, e); + } catch (RestClientException e) { + LOGGER.error("Schema Registry error for subject '{}'", subject, e); + throw new ProcessorInitializationException("Schema Registry error for subject: " + subject, e); + } + } + return schemaType; + } + + private String getSubjectName(String topic) { + return topic + VALUE_SUFFIX; + } + + private String getSubjectName(String topic, boolean isKey) { + return topic + (isKey ? KEY_SUFFIX : VALUE_SUFFIX); + } + + private Converter createConverterForFormat(String format, boolean isKey) { + LOGGER.info("Creating new converter for format: {}", format); + SchemaFormat schemaFormat = SchemaFormat.fromString(format); + switch (schemaFormat) { + case AVRO: + return new AvroRegistryConverter(client, schemaRegistryUrl, isKey); + case PROTOBUF: + return new ProtobufRegistryConverter(client, schemaRegistryUrl, isKey); + default: + LOGGER.error("Unsupported schema format '{}'", format); + throw new ProcessorInitializationException("Unsupported schema format: " + format); + } + } + +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/LazyConverter.java b/core/src/main/java/kafka/automq/table/process/convert/LazyConverter.java new file mode 100644 index 0000000000..4aad9eca17 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/LazyConverter.java @@ -0,0 +1,80 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.convert; + +import kafka.automq.table.process.ConversionResult; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.exception.ConverterException; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.nio.ByteBuffer; +import java.util.function.Supplier; + +/** + * A decorator for a {@link Converter} that delays its initialization until it is first used. + * + *

This is useful for deferring expensive initialization logic, such as fetching + * schema metadata over the network, until a record from the corresponding topic + * is actually being processed.

+ * + *

This class is thread-safe.

+ */ +public class LazyConverter implements Converter { + private static final Logger log = LoggerFactory.getLogger(LazyConverter.class); + + private final Supplier converterSupplier; + private volatile Converter delegate; + + /** + * Creates a new LazyConverter. + * + * @param converterSupplier a supplier that provides the actual converter instance when called. + */ + public LazyConverter(Supplier converterSupplier) { + this.converterSupplier = converterSupplier; + } + + /** + * Gets the delegate converter, initializing it if necessary. + * Uses double-checked locking for thread-safe lazy initialization. + */ + private Converter getDelegate() { + if (delegate == null) { + synchronized (this) { + if (delegate == null) { + Converter localDelegate = converterSupplier.get(); + if (localDelegate == null) { + throw new IllegalStateException("Converter supplier returned null"); + } + log.info("Successfully initialized delegate converter: {}", localDelegate.getClass().getName()); + delegate = localDelegate; + } + } + } + return delegate; + } + + @Override + public ConversionResult convert(String topic, ByteBuffer buffer) throws ConverterException { + return getDelegate().convert(topic, buffer); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/LogicalMapProtobufData.java b/core/src/main/java/kafka/automq/table/process/convert/LogicalMapProtobufData.java new file mode 100644 index 0000000000..5206a3d2cd --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/LogicalMapProtobufData.java @@ -0,0 +1,77 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import com.google.protobuf.Descriptors; + +import org.apache.avro.Schema; +import org.apache.avro.protobuf.ProtobufData; +import org.apache.iceberg.avro.CodecSetup; + +import java.util.Arrays; + +/** + * ProtobufData extension that annotates protobuf map fields with Iceberg's LogicalMap logical type so that + * downstream Avro{@literal >}Iceberg conversion keeps them as MAP instead of generic {@literal ARRAY>}. + */ +public class LogicalMapProtobufData extends ProtobufData { + private static final LogicalMapProtobufData INSTANCE = new LogicalMapProtobufData(); + private static final Schema NULL = Schema.create(Schema.Type.NULL); + + public static LogicalMapProtobufData get() { + return INSTANCE; + } + + @Override + public Schema getSchema(Descriptors.FieldDescriptor f) { + Schema schema = super.getSchema(f); + if (f.isMapField()) { + Schema nonNull = resolveNonNull(schema); + // protobuf maps are materialized as ARRAY in Avro + if (nonNull != null && nonNull.getType() == Schema.Type.ARRAY) { + // set logicalType property; LogicalTypes is registered in CodecSetup + CodecSetup.getLogicalMap().addToSchema(nonNull); + } + } else if (f.isOptional() && !f.isRepeated() && f.getContainingOneof() == null + && schema.getType() != Schema.Type.UNION) { + // Proto3 optional scalars/messages: wrap as union(type, null) so the protobuf default (typically non-null) + // remains valid (Avro default must match the first branch). + schema = Schema.createUnion(Arrays.asList(schema, NULL)); + } else if (f.getContainingOneof() != null && !f.isRepeated() && schema.getType() != Schema.Type.UNION) { + // oneof fields: wrap as union(type, null) so that non-set fields can be represented as null + schema = Schema.createUnion(Arrays.asList(schema, NULL)); + } + return schema; + } + + private Schema resolveNonNull(Schema schema) { + if (schema == null) { + return null; + } + if (schema.getType() == Schema.Type.UNION) { + for (Schema member : schema.getTypes()) { + if (member.getType() != Schema.Type.NULL) { + return member; + } + } + return null; + } + return schema; + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/ProtoToAvroConverter.java b/core/src/main/java/kafka/automq/table/process/convert/ProtoToAvroConverter.java new file mode 100644 index 0000000000..e6f0ec2130 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/ProtoToAvroConverter.java @@ -0,0 +1,196 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.convert; + +import kafka.automq.table.process.exception.ConverterException; + +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import com.google.protobuf.Message; +import com.google.protobuf.Timestamp; + +import org.apache.avro.Conversion; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.protobuf.ProtoConversions; +import org.apache.avro.protobuf.ProtobufData; + +import java.nio.ByteBuffer; +import java.util.List; + +public class ProtoToAvroConverter { + + private static final ProtobufData DATA = initProtobufData(); + + private static ProtobufData initProtobufData() { + ProtobufData protobufData = LogicalMapProtobufData.get(); + protobufData.addLogicalTypeConversion(new ProtoConversions.TimestampMicrosConversion()); + return protobufData; + } + + public static GenericRecord convert(Message protoMessage, Schema schema) { + try { + Schema nonNull = resolveNonNullSchema(schema); + return convertRecord(protoMessage, nonNull, DATA); + } catch (Exception e) { + throw new ConverterException("Proto to Avro conversion failed", e); + } + } + + private static Object convert(Message protoMessage, Schema schema, ProtobufData protobufData) { + Conversion conversion = getConversion(protoMessage.getDescriptorForType(), protobufData); + if (conversion instanceof ProtoConversions.TimestampMicrosConversion) { + ProtoConversions.TimestampMicrosConversion timestampConversion = (ProtoConversions.TimestampMicrosConversion) conversion; + Timestamp.Builder builder = Timestamp.newBuilder(); + Timestamp.getDescriptor().getFields().forEach(field -> { + Descriptors.FieldDescriptor protoField = protoMessage.getDescriptorForType().findFieldByName(field.getName()); + if (protoField != null && protoMessage.hasField(protoField)) { + builder.setField(field, protoMessage.getField(protoField)); + } + }); + return timestampConversion.toLong(builder.build(), schema, null); + } + + Schema nonNull = resolveNonNullSchema(schema); + if (nonNull.getType() == Schema.Type.RECORD) { + return convertRecord(protoMessage, nonNull, protobufData); + } + return null; + } + + private static Conversion getConversion(Descriptors.Descriptor descriptor, ProtobufData protobufData) { + String namespace = protobufData.getNamespace(descriptor.getFile(), descriptor.getContainingType()); + String name = descriptor.getName(); + if ("com.google.protobuf".equals(namespace) && "Timestamp".equals(name)) { + return new ProtoConversions.TimestampMicrosConversion(); + } + return null; + } + + private static GenericRecord convertRecord(Message protoMessage, Schema recordSchema, ProtobufData protobufData) { + GenericRecord record = new GenericData.Record(recordSchema); + Descriptors.Descriptor descriptor = protoMessage.getDescriptorForType(); + + for (Schema.Field field : recordSchema.getFields()) { + String fieldName = field.name(); + Descriptors.FieldDescriptor protoField = descriptor.findFieldByName(fieldName); + if (protoField == null) { + continue; + } + + boolean hasPresence = protoField.hasPresence() || protoField.getContainingOneof() != null; + if (!protoField.isRepeated() && hasPresence && !protoMessage.hasField(protoField)) { + if (allowsNull(field.schema())) { + record.put(fieldName, null); + } + continue; + } + + Object value = protoMessage.getField(protoField); + Object convertedValue = convertValue(value, protoField, field.schema(), protobufData); + record.put(fieldName, convertedValue); + } + return record; + } + + private static Object convertValue(Object value, Descriptors.FieldDescriptor fieldDesc, Schema avroSchema, + ProtobufData protobufData) { + if (value == null) { + return null; + } + + Schema nonNullSchema = resolveNonNullSchema(avroSchema); + + if (fieldDesc.isRepeated() && value instanceof List) { + List protoList = (List) value; + GenericData.Array avroArray = new GenericData.Array<>(protoList.size(), nonNullSchema); + Schema elementSchema = nonNullSchema.getElementType(); + for (Object item : protoList) { + avroArray.add(convertSingleValue(item, elementSchema, protobufData)); + } + return avroArray; + } + + return convertSingleValue(value, nonNullSchema, protobufData); + } + + private static Object convertSingleValue(Object value, Schema avroSchema, ProtobufData protobufData) { + if (value instanceof Message) { + return convert((Message) value, avroSchema, protobufData); + } else if (value instanceof ByteString) { + return ((ByteString) value).asReadOnlyByteBuffer(); + } else if (value instanceof Enum) { + return value.toString(); + } else if (value instanceof List) { + throw new ConverterException("Unexpected list type found; repeated fields should have been handled in convertValue"); + } + + return convertPrimitive(value, avroSchema); + } + + private static Object convertPrimitive(Object value, Schema schema) { + Schema.Type type = schema.getType(); + switch (type) { + case INT: + return ((Number) value).intValue(); + case LONG: + return ((Number) value).longValue(); + case FLOAT: + return ((Number) value).floatValue(); + case DOUBLE: + return ((Number) value).doubleValue(); + case BOOLEAN: + return (Boolean) value; + case BYTES: + if (value instanceof byte[]) { + return ByteBuffer.wrap((byte[]) value); + } + return value; + default: + return value; + } + } + + private static Schema resolveNonNullSchema(Schema schema) { + if (schema.getType() == Schema.Type.UNION) { + for (Schema type : schema.getTypes()) { + if (type.getType() != Schema.Type.NULL) { + return type; + } + } + } + return schema; + } + + private static boolean allowsNull(Schema schema) { + if (schema.getType() == Schema.Type.NULL) { + return true; + } + if (schema.getType() == Schema.Type.UNION) { + for (Schema type : schema.getTypes()) { + if (type.getType() == Schema.Type.NULL) { + return true; + } + } + } + return false; + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/ProtobufRegistryConverter.java b/core/src/main/java/kafka/automq/table/process/convert/ProtobufRegistryConverter.java new file mode 100644 index 0000000000..ffa01b881f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/ProtobufRegistryConverter.java @@ -0,0 +1,89 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import kafka.automq.table.deserializer.SchemaResolutionResolver; +import kafka.automq.table.deserializer.proto.CustomKafkaProtobufDeserializer; +import kafka.automq.table.deserializer.proto.HeaderBasedSchemaResolutionResolver; +import kafka.automq.table.process.ConversionResult; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.exception.ConverterException; +import kafka.automq.table.process.exception.InvalidDataException; + +import org.apache.kafka.common.serialization.Deserializer; + +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import com.google.protobuf.Message; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.protobuf.ProtoConversions; +import org.apache.avro.protobuf.ProtobufData; + +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.Map; + +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; + +public class ProtobufRegistryConverter implements Converter { + private final Deserializer deserializer; + private final SchemaResolutionResolver resolver; + + private final Cache avroSchemaCache = CacheBuilder.newBuilder() + .expireAfterAccess(Duration.ofMinutes(20)) + .maximumSize(10000) + .build(); + + public ProtobufRegistryConverter(Deserializer deserializer) { + this.deserializer = deserializer; + this.resolver = new HeaderBasedSchemaResolutionResolver(); + } + + public ProtobufRegistryConverter(SchemaRegistryClient client, String registryUrl, boolean isKey) { + this(client, registryUrl, new HeaderBasedSchemaResolutionResolver(), isKey); + } + + public ProtobufRegistryConverter(SchemaRegistryClient client, String registryUrl, SchemaResolutionResolver resolver, boolean isKey) { + this.resolver = resolver; + this.deserializer = new CustomKafkaProtobufDeserializer<>(client, resolver); + // Configure the deserializer immediately upon creation + Map configs = Map.of("schema.registry.url", registryUrl); + deserializer.configure(configs, isKey); + } + + @Override + public ConversionResult convert(String topic, ByteBuffer buffer) throws ConverterException { + if (buffer == null) { + throw new InvalidDataException("buffer is null"); + } + int schemaId = resolver.getSchemaId(topic, buffer); + Message protoMessage = deserializer.deserialize(topic, null, buffer); + Schema schema = avroSchemaCache.getIfPresent(schemaId); + if (schema == null) { + ProtobufData protobufData = LogicalMapProtobufData.get(); + protobufData.addLogicalTypeConversion(new ProtoConversions.TimestampMicrosConversion()); + schema = protobufData.getSchema(protoMessage.getDescriptorForType()); + avroSchemaCache.put(schemaId, schema); + } + GenericRecord convert = ProtoToAvroConverter.convert(protoMessage, schema); + return new ConversionResult(convert, String.valueOf(schemaId)); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/RawConverter.java b/core/src/main/java/kafka/automq/table/process/convert/RawConverter.java new file mode 100644 index 0000000000..f358426272 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/RawConverter.java @@ -0,0 +1,46 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import kafka.automq.table.process.ConversionResult; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.exception.ConverterException; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.SchemaNormalization; + +import java.nio.ByteBuffer; + +public class RawConverter implements Converter { + public static final RawConverter INSTANCE = new RawConverter(); + + private static final Schema SCHEMA = SchemaBuilder.builder().bytesType(); + private static final String SCHEMA_IDENTITY = Long.toUnsignedString(SchemaNormalization.parsingFingerprint64(SCHEMA)); + + @Override + public ConversionResult convert(String topic, ByteBuffer buffer) throws ConverterException { + if (buffer == null) { + return new ConversionResult(new byte[0], SCHEMA, SCHEMA_IDENTITY); + } + byte[] bytes = new byte[buffer.remaining()]; + buffer.slice().get(bytes); + return new ConversionResult(bytes, SCHEMA, SCHEMA_IDENTITY); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/convert/StringConverter.java b/core/src/main/java/kafka/automq/table/process/convert/StringConverter.java new file mode 100644 index 0000000000..f2e6081359 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/convert/StringConverter.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import kafka.automq.table.process.ConversionResult; +import kafka.automq.table.process.Converter; +import kafka.automq.table.process.exception.ConverterException; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.SchemaNormalization; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; + +public class StringConverter implements Converter { + public static final StringConverter INSTANCE = new StringConverter(); + + private static final Schema SCHEMA = SchemaBuilder.builder().stringType(); + private static final String SCHEMA_IDENTITY = Long.toUnsignedString(SchemaNormalization.parsingFingerprint64(SCHEMA)); + + @Override + public ConversionResult convert(String topic, ByteBuffer buffer) throws ConverterException { + if (buffer == null) { + return new ConversionResult("", SCHEMA, SCHEMA_IDENTITY); + } + byte[] bytes = new byte[buffer.remaining()]; + buffer.slice().get(bytes); + return new ConversionResult(new String(bytes, StandardCharsets.UTF_8), SCHEMA, SCHEMA_IDENTITY); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/exception/ConverterException.java b/core/src/main/java/kafka/automq/table/process/exception/ConverterException.java new file mode 100644 index 0000000000..b6b5280ec4 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/exception/ConverterException.java @@ -0,0 +1,20 @@ +package kafka.automq.table.process.exception; + +public class ConverterException extends RuntimeException { + + public ConverterException() { + super(); + } + + public ConverterException(String message) { + super(message); + } + + public ConverterException(String message, Throwable cause) { + super(message, cause); + } + + public ConverterException(Throwable cause) { + super(cause); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/exception/InvalidDataException.java b/core/src/main/java/kafka/automq/table/process/exception/InvalidDataException.java new file mode 100644 index 0000000000..00eac41413 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/exception/InvalidDataException.java @@ -0,0 +1,39 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.exception; + +public class InvalidDataException extends RuntimeException { + private static final long serialVersionUID = 4029025366392702726L; + + public InvalidDataException() { + } + + public InvalidDataException(String msg) { + super(msg); + } + + public InvalidDataException(String msg, Throwable cause) { + super(msg, cause); + } + + public InvalidDataException(Throwable cause) { + super(cause); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/exception/ProcessorInitializationException.java b/core/src/main/java/kafka/automq/table/process/exception/ProcessorInitializationException.java new file mode 100644 index 0000000000..ac6b5281d1 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/exception/ProcessorInitializationException.java @@ -0,0 +1,45 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.exception; + +/** + * Unchecked exception for fatal, system-level errors during processor initialization. + * + *

This exception represents unrecoverable errors that occur during the setup + * and configuration phase, such as: + *

    + *
  • Failure to connect to a required external service (e.g., Schema Registry)
    • + *
  • Invalid or inconsistent configuration that prevents component creation
    • + *
+ * + *

Unlike {@link RecordProcessorException}, this is a {@link RuntimeException} + * because initialization errors are typically programming or deployment errors + * that should not be caught and handled at runtime. They should cause the + * application to fail fast.

+ */ +public class ProcessorInitializationException extends RuntimeException { + + public ProcessorInitializationException(String message) { + super(message); + } + + public ProcessorInitializationException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/exception/RecordProcessorException.java b/core/src/main/java/kafka/automq/table/process/exception/RecordProcessorException.java new file mode 100644 index 0000000000..487cfc9f2b --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/exception/RecordProcessorException.java @@ -0,0 +1,57 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.exception; + +public class RecordProcessorException extends Exception { + /** + * Constructs a new record processor exception with no detail message. + */ + public RecordProcessorException() { + super(); + } + + /** + * Constructs a new record processor exception with the specified detail message. + * + * @param message the detail message explaining the cause of the exception + */ + public RecordProcessorException(String message) { + super(message); + } + + /** + * Constructs a new record processor exception with the specified detail message and cause. + * + * @param message the detail message explaining the cause of the exception + * @param cause the cause of the exception (which is saved for later retrieval) + */ + public RecordProcessorException(String message, Throwable cause) { + super(message, cause); + } + + /** + * Constructs a new record processor exception with the specified cause. + * + * @param cause the cause of the exception (which is saved for later retrieval) + */ + public RecordProcessorException(Throwable cause) { + super(cause); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/exception/SchemaRegistrySystemException.java b/core/src/main/java/kafka/automq/table/process/exception/SchemaRegistrySystemException.java new file mode 100644 index 0000000000..86e1c3eca0 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/exception/SchemaRegistrySystemException.java @@ -0,0 +1,118 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.exception; + + +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; + +/** + * Exception representing system-level errors when interacting with the Schema Registry. + * These are critical errors that typically cannot be tolerated even with ErrorsTolerance=ALL + * configuration, as they represent system infrastructure issues rather than data problems. + */ +public class SchemaRegistrySystemException extends RuntimeException { + + private final ErrorType errorType; + private final int statusCode; + + /** + * Classification of schema registry system errors + */ + public enum ErrorType { + /** + * Authentication errors (401) + */ + AUTHENTICATION_ERROR, + + /** + * Authorization errors (403) + */ + AUTHORIZATION_ERROR, + + /** + * Rate limiting errors (429) + */ + RATE_LIMIT_ERROR, + + /** + * Temporary service unavailable errors (408, 503, 504) + */ + SERVICE_UNAVAILABLE_ERROR, + + /** + * Gateway errors (502) indicating upstream service issues + */ + GATEWAY_ERROR, + + /** + * Other unexpected system errors + */ + UNKNOWN_SYSTEM_ERROR + } + + public SchemaRegistrySystemException(String message, Throwable cause, ErrorType errorType, int statusCode) { + super(message, cause); + this.errorType = errorType; + this.statusCode = statusCode; + } + + public ErrorType getErrorType() { + return errorType; + } + + public int getStatusCode() { + return statusCode; + } + + + + // io.confluent.kafka.serializers.AbstractKafkaSchemaSerDe#toKafkaException + public static SchemaRegistrySystemException fromStatusCode(RestClientException exception, String recordContext) { + int status = exception.getStatus(); + + if (status == 401) { + return new SchemaRegistrySystemException( + "Authentication error when accessing schema registry for record: " + recordContext, + exception, ErrorType.AUTHENTICATION_ERROR, status); + } else if (status == 403) { + return new SchemaRegistrySystemException( + "Authorization error when accessing schema registry for record: " + recordContext, + exception, ErrorType.AUTHORIZATION_ERROR, status); + } else if (status == 429) { // Too Many Requests + return new SchemaRegistrySystemException( + "Rate limit exceeded when accessing schema registry for record: " + recordContext, + exception, ErrorType.RATE_LIMIT_ERROR, status); + } else if (status == 408 // Request Timeout + || status == 503 // Service Unavailable + || status == 504) { // Gateway Timeout + return new SchemaRegistrySystemException( + "Service unavailable or timeout when accessing schema registry for record: " + recordContext, + exception, ErrorType.SERVICE_UNAVAILABLE_ERROR, status); + } else if (status == 502) { // Bad Gateway + return new SchemaRegistrySystemException( + "Bad gateway error when accessing schema registry for record: " + recordContext, + exception, ErrorType.GATEWAY_ERROR, status); + } else { + return new SchemaRegistrySystemException( + "Unexpected schema registry error (HTTP " + status + ") for record: " + recordContext, + exception, ErrorType.UNKNOWN_SYSTEM_ERROR, status); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/exception/TransformException.java b/core/src/main/java/kafka/automq/table/process/exception/TransformException.java new file mode 100644 index 0000000000..2dc543ecf7 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/exception/TransformException.java @@ -0,0 +1,39 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.exception; + +public class TransformException extends RuntimeException { + + public TransformException() { + super(); + } + + public TransformException(String message) { + super(message); + } + + public TransformException(String message, Throwable cause) { + super(message, cause); + } + + public TransformException(Throwable cause) { + super(cause); + } +} diff --git a/core/src/main/java/kafka/automq/table/process/transform/DebeziumUnwrapTransform.java b/core/src/main/java/kafka/automq/table/process/transform/DebeziumUnwrapTransform.java new file mode 100644 index 0000000000..a4e2e5845d --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/transform/DebeziumUnwrapTransform.java @@ -0,0 +1,323 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.transform; + +import kafka.automq.table.process.Transform; +import kafka.automq.table.process.TransformContext; +import kafka.automq.table.process.exception.TransformException; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +import io.confluent.kafka.schemaregistry.utils.BoundedConcurrentHashMap; + +/** + * Transform for unwrapping Debezium CDC formatted records. + * + *

Extracts business data from Debezium CDC envelope format. Handles create, + * update, and delete operations, adding unified CDC metadata.

+ * + *
    + *
  • CREATE/READ: uses 'after' field
    • + *
  • UPDATE: uses 'after' field
    • + *
  • DELETE: uses 'before' field
    • + *
+ * + *

Adds metadata: operation type, timestamp, offset, source table.

+ */ +public class DebeziumUnwrapTransform implements Transform { + + public static final DebeziumUnwrapTransform INSTANCE = new DebeziumUnwrapTransform(); + + private static final Logger log = LoggerFactory.getLogger(DebeziumUnwrapTransform.class); + + // Debezium standard field names + private static final String FIELD_BEFORE = "before"; + private static final String FIELD_AFTER = "after"; + private static final String FIELD_OP = "op"; + private static final String FIELD_SOURCE = "source"; + private static final String FIELD_TS_MS = "ts_ms"; + + // Debezium operation types + private static final String OP_CREATE = "c"; + private static final String OP_UPDATE = "u"; + private static final String OP_DELETE = "d"; + private static final String OP_READ = "r"; + + //CDC field names + private static final String CDC_RECORD_NAME = "_cdc"; + private static final String CDC_FIELD_OP = "op"; + private static final String CDC_FIELD_TS = "ts"; + private static final String CDC_FIELD_OFFSET = "offset"; + private static final String CDC_FIELD_SOURCE = "source"; + + private static final Schema CDC_SCHEMA = SchemaBuilder.record(CDC_RECORD_NAME) + .fields() + .optionalString(CDC_FIELD_OP) + .optionalLong(CDC_FIELD_TS) + .optionalLong(CDC_FIELD_OFFSET) + .optionalString(CDC_FIELD_SOURCE) + .endRecord(); + + // Cache enriched schemas keyed by base schema fingerprint (bounded, concurrent) + private static final int ENRICHED_SCHEMA_CACHE_MAX = 1024; + private final Map enrichedSchemaCache = new BoundedConcurrentHashMap<>(ENRICHED_SCHEMA_CACHE_MAX); + + @Override + public void configure(Map configs) { + // ignore + } + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) throws TransformException { + Objects.requireNonNull(record, "Input record cannot be null"); + + try { + // If it's not a Debezium record, throw an exception. + if (!isDebeziumRecord(record)) { + throw new TransformException("Record is not in a recognizable Debezium format."); + } + + // Extract operation type + String operation = getStringValue(record, FIELD_OP); + if (operation == null) { + throw new TransformException("Invalid Debezium record: missing required field '" + FIELD_OP + "'"); + } + + // Extract business data based on operation type + GenericRecord businessData = extractBusinessData(record, operation); + if (businessData == null) { + throw new TransformException("Invalid Debezium record: no extractable data for operation '" + operation + "'"); + } + + // Enrich with metadata + return enrichWithMetadata(businessData, record, operation, context); + + } catch (TransformException e) { + throw e; + } catch (Exception e) { + throw new TransformException("Failed to process Debezium record", e); + } + } + + private boolean isDebeziumRecord(GenericRecord record) { + if (record == null) { + return false; + } + Schema schema = unwrapSchema(record.getSchema()); + if (schema == null) { + return false; + } + return schema.getField(FIELD_OP) != null && + (schema.getField(FIELD_BEFORE) != null || schema.getField(FIELD_AFTER) != null); + } + + private GenericRecord extractBusinessData(GenericRecord record, String operation) throws TransformException { + switch (operation) { + case OP_CREATE: + case OP_READ: + // INSERT and READ operations use 'after' field + return getRecordValue(record, FIELD_AFTER); + + case OP_UPDATE: + // UPDATE operations must have 'after' field + GenericRecord after = getRecordValue(record, FIELD_AFTER); + if (after == null) { + throw new TransformException("Invalid UPDATE record: missing required 'after' data"); + } + return after; + + case OP_DELETE: + // DELETE operations use 'before' field + GenericRecord beforeDelete = getRecordValue(record, FIELD_BEFORE); + if (beforeDelete == null) { + throw new TransformException("Invalid DELETE record: missing required 'before' data"); + } + return beforeDelete; + + default: + log.warn("Unknown Debezium operation type: {}. Attempting to use 'after' data", operation); + GenericRecord fallback = getRecordValue(record, FIELD_AFTER); + if (fallback == null) { + throw new TransformException("Unsupported operation '" + operation + "' with no usable data"); + } + return fallback; + } + } + + + private GenericRecord enrichWithMetadata(GenericRecord businessData, + GenericRecord debeziumRecord, + String operation, + TransformContext context) throws TransformException { + try { + Schema schemaWithMetadata = createSchemaWithMetadata(businessData.getSchema()); + + // Build the enriched record using GenericData.Record to reduce allocations + GenericData.Record result = new GenericData.Record(schemaWithMetadata); + for (Schema.Field field : businessData.getSchema().getFields()) { + result.put(field.name(), businessData.get(field.name())); + } + + GenericData.Record cdc = new GenericData.Record(CDC_SCHEMA); + cdc.put(CDC_FIELD_OP, mapOperation(operation)); + + Object tsMs = debeziumRecord.get(FIELD_TS_MS); + if (tsMs instanceof Long) { + cdc.put(CDC_FIELD_TS, tsMs); + } + + cdc.put(CDC_FIELD_OFFSET, context.getKafkaRecord().offset()); + + GenericRecord source = getRecordValue(debeziumRecord, FIELD_SOURCE); + if (source != null) { + String schemaName = null; + if (source.hasField("schema")) { + schemaName = getStringValue(source, "schema"); + } + String db = (schemaName == null) ? getStringValue(source, "db") : schemaName; + String table = getStringValue(source, "table"); + if (db != null && table != null) { + cdc.put(CDC_FIELD_SOURCE, db + "." + table); + } + } + + result.put(CDC_RECORD_NAME, cdc); + return result; + + } catch (Exception e) { + throw new TransformException("Failed to enrich record with Debezium metadata:" + e.getMessage(), e); + } + } + + private String mapOperation(String originalOp) { + switch (originalOp) { + case "u": + return "U"; + case "d": + return "D"; + default: + // Debezium ops "c", "r", and any others + return "I"; + } + } + + private Schema createSchemaWithMetadata(Schema originalSchema) { + SchemaKey schemaKey = new SchemaKey(originalSchema); + return enrichedSchemaCache.computeIfAbsent(schemaKey, k -> { + List enhancedFields = new ArrayList<>(); + for (Schema.Field field : originalSchema.getFields()) { + enhancedFields.add(new Schema.Field(field, field.schema())); + } + enhancedFields.add(new Schema.Field(CDC_RECORD_NAME, CDC_SCHEMA, "CDC metadata", null)); + + String enhancedName = originalSchema.getName() != null ? + originalSchema.getName() + "_cdc_enriched" : "enriched_record"; + + return Schema.createRecord( + enhancedName, + "Record enriched with CDC metadata", + originalSchema.getNamespace(), + false, + enhancedFields + ); + }); + } + + + private GenericRecord getRecordValue(GenericRecord record, String fieldName) { + Object value = record.get(fieldName); + return (value instanceof GenericRecord) ? (GenericRecord) value : null; + } + private String getStringValue(GenericRecord record, String fieldName) { + Object value = record.get(fieldName); + return (value != null) ? value.toString() : null; + } + + private Schema unwrapSchema(Schema schema) { + if (schema == null) { + return null; + } + if (schema.getType() == Schema.Type.UNION) { + return schema.getTypes().stream() + .filter(s -> s.getType() == Schema.Type.RECORD) + .findFirst() + .orElse(null); + } + return schema.getType() == Schema.Type.RECORD ? schema : null; + } + + @Override + public String getName() { + return "DebeziumUnwrap"; + } + + private static final class SchemaKey { + private final Schema schema; + private final int hashCode; + private volatile long fingerprint; + private volatile boolean fingerprintComputed; + + private SchemaKey(Schema schema) { + this.schema = Objects.requireNonNull(schema, "schema"); + this.hashCode = schema.hashCode(); + } + + @Override + public int hashCode() { + return hashCode; + } + + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof SchemaKey)) { + return false; + } + SchemaKey other = (SchemaKey) obj; + if (hashCode != other.hashCode) { + return false; + } + if (schema == other.schema) { + return true; + } + return fingerprint() == other.fingerprint(); + } + + private long fingerprint() { + if (!fingerprintComputed) { + fingerprint = org.apache.avro.SchemaNormalization.parsingFingerprint64(schema); + fingerprintComputed = true; + } + return fingerprint; + } + } +} diff --git a/core/src/main/java/kafka/automq/table/process/transform/FlattenTransform.java b/core/src/main/java/kafka/automq/table/process/transform/FlattenTransform.java new file mode 100644 index 0000000000..d15ebb7bc9 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/transform/FlattenTransform.java @@ -0,0 +1,60 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.transform; + +import kafka.automq.table.process.RecordAssembler; +import kafka.automq.table.process.Transform; +import kafka.automq.table.process.TransformContext; +import kafka.automq.table.process.exception.TransformException; + +import org.apache.avro.generic.GenericRecord; + +import java.util.Map; + +/** + * A transform to unwrap a record from a standard {@code ValueRecord} container. + */ +public class FlattenTransform implements Transform { + + public static final FlattenTransform INSTANCE = new FlattenTransform(); + + @Override + public void configure(Map configs) { + // No configuration needed for this transform. + } + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) throws TransformException { + if (record == null || !record.hasField(RecordAssembler.KAFKA_VALUE_FIELD)) { + throw new TransformException("Record is null or has no value field"); + } + Object value = record.get(RecordAssembler.KAFKA_VALUE_FIELD); + if (value instanceof GenericRecord) { + return (GenericRecord) value; + } else { + throw new TransformException("value field is not a GenericRecord"); + } + } + + @Override + public String getName() { + return "Flatten"; + } +} diff --git a/core/src/main/java/kafka/automq/table/process/transform/SchemalessTransform.java b/core/src/main/java/kafka/automq/table/process/transform/SchemalessTransform.java new file mode 100644 index 0000000000..708d318ca2 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/process/transform/SchemalessTransform.java @@ -0,0 +1,105 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.transform; + +import kafka.automq.table.process.Transform; +import kafka.automq.table.process.TransformContext; +import kafka.automq.table.process.exception.TransformException; + +import org.apache.kafka.common.record.Record; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Map; + +/** + * SchemalessTransform provides a simplified transformation approach that ignores the structured data + * from valueConverter. It extracts the raw bytes of key, value, and timestamp from the original Kafka Record + * and assembles them into a flat GenericRecord. + *

+ * This transform is designed for compatibility with legacy "schemaless" configuration. + */ +public class SchemalessTransform implements Transform { + + private static final String TRANSFORM_NAME = "schemaless"; + private static final String FIELD_KEY = "key"; + private static final String FIELD_VALUE = "value"; + private static final String FIELD_TIMESTAMP = "timestamp"; + + public static final Schema SCHEMALESS_SCHEMA = SchemaBuilder + .record("SchemalessRecord") + .namespace("kafka.automq.table.process.transform") + .doc("A simple record containing raw key, value, and timestamp.") + .fields() + .name(FIELD_KEY).doc("Original record key as string") + .type().unionOf().nullType().and().stringType().endUnion() + .nullDefault() + .name(FIELD_VALUE).doc("Original record value as string") + .type().unionOf().nullType().and().stringType().endUnion() + .nullDefault() + .name(FIELD_TIMESTAMP).doc("Record timestamp") + .type().longType() + .longDefault(0L) + .endRecord(); + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) throws TransformException { + // Ignore the input record, we only use the context + try { + Record kafkaRecord = context.getKafkaRecord(); + + // Create a new record with the predefined schema + GenericRecord schemalessRecord = new GenericData.Record(SCHEMALESS_SCHEMA); + + // Set fields using constants + schemalessRecord.put(FIELD_KEY, kafkaRecord.hasKey() ? buf2String(kafkaRecord.key()) : null); + schemalessRecord.put(FIELD_VALUE, kafkaRecord.hasValue() ? buf2String(kafkaRecord.value()) : null); + schemalessRecord.put(FIELD_TIMESTAMP, kafkaRecord.timestamp()); + + return schemalessRecord; + } catch (Exception e) { + throw new TransformException("Failed to process record in SchemalessTransform", e); + } + } + + @Override + public void configure(Map configs) { + // Nothing to configure for this transform + } + + @Override + public String getName() { + return TRANSFORM_NAME; + } + + private String buf2String(ByteBuffer buffer) { + if (buffer == null) { + return ""; + } + byte[] bytes = new byte[buffer.remaining()]; + buffer.slice().get(bytes); + return new String(bytes, StandardCharsets.UTF_8); + } +} diff --git a/core/src/main/java/kafka/automq/table/utils/PartitionUtil.java b/core/src/main/java/kafka/automq/table/utils/PartitionUtil.java new file mode 100644 index 0000000000..ae85113316 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/utils/PartitionUtil.java @@ -0,0 +1,201 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.utils; + +import org.apache.kafka.server.common.automq.TableTopicConfigValidator; + +import org.apache.commons.lang3.tuple.Pair; +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.UpdatePartitionSpec; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.expressions.NamedReference; +import org.apache.iceberg.expressions.Term; +import org.apache.iceberg.expressions.UnboundTransform; +import org.apache.iceberg.transforms.Transforms; +import org.apache.iceberg.types.Types; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import static org.apache.iceberg.expressions.Expressions.bucket; +import static org.apache.kafka.server.common.automq.TableTopicConfigValidator.PartitionValidator.transformArgPair; + +@SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"}) +public class PartitionUtil { + public static final Pattern TRANSFORM_REGEX = Pattern.compile("(\\w+)\\((.+)\\)"); + + public static List parsePartitionBy(String str) { + return TableTopicConfigValidator.PartitionValidator.parsePartitionBy(str); + } + + public static PartitionSpec buildPartitionSpec(List partitions, Schema schema) { + if (partitions.isEmpty()) { + return PartitionSpec.unpartitioned(); + } + PartitionSpec.Builder specBuilder = PartitionSpec.builderFor(schema); + partitions.forEach( + partitionField -> { + Matcher matcher = TRANSFORM_REGEX.matcher(partitionField); + if (matcher.matches()) { + String transform = matcher.group(1); + switch (transform) { + case "year": + specBuilder.year(matcher.group(2)); + break; + case "month": + specBuilder.month(matcher.group(2)); + break; + case "day": + specBuilder.day(matcher.group(2)); + break; + case "hour": + specBuilder.hour(matcher.group(2)); + break; + case "bucket": { + Pair args = transformArgPair(matcher.group(2)); + specBuilder.bucket(args.getLeft(), args.getRight()); + break; + } + case "truncate": { + Pair args = transformArgPair(matcher.group(2)); + specBuilder.truncate(args.getLeft(), args.getRight()); + break; + } + default: + throw new UnsupportedOperationException("Unsupported transform: " + transform); + } + } else { + specBuilder.identity(partitionField); + } + }); + return specBuilder.build(); + } + + public static boolean evolve(List newPartitions, Table table) { + PartitionSpec spec = table.spec(); + if (newPartitions.isEmpty() && spec.isUnpartitioned()) { + return false; + } + Map id2field = new HashMap<>(); + Set newPartitionFieldIdSet = new HashSet<>(); + spec.fields().forEach(f -> id2field.put(f.sourceId(), f)); + + Schema tableSchema = table.schema(); + UpdatePartitionSpec updateSpec = table.updateSpec(); + int changeCount = 0; + // add/replace partition + for (String newPartition : newPartitions) { + Matcher matcher = TRANSFORM_REGEX.matcher(newPartition); + String transformer; + String fieldName; + if (matcher.matches()) { + transformer = matcher.group(1); + if ("bucket".equals(transformer) || "truncate".equals(transformer)) { + Pair args = transformArgPair(matcher.group(2)); + fieldName = args.getLeft(); + } else { + fieldName = matcher.group(2); + } + } else { + transformer = "identity"; + fieldName = newPartition; + } + Types.NestedField nestedField = tableSchema.findField(fieldName); + if (nestedField == null) { + continue; + } + newPartitionFieldIdSet.add(nestedField.fieldId()); + switch (transformer) { + case "year": { + changeCount += addOrUpdate(nestedField, Expressions.year(fieldName), updateSpec, id2field); + break; + } + case "month": { + changeCount += addOrUpdate(nestedField, Expressions.month(fieldName), updateSpec, id2field); + break; + } + case "day": { + changeCount += addOrUpdate(nestedField, Expressions.day(fieldName), updateSpec, id2field); + break; + } + case "hour": { + changeCount += addOrUpdate(nestedField, Expressions.hour(fieldName), updateSpec, id2field); + break; + } + case "bucket": { + Pair args = transformArgPair(matcher.group(2)); + changeCount += addOrUpdate(nestedField, bucket(args.getLeft(), args.getRight()), updateSpec, id2field); + break; + } + case "truncate": { + Pair args = transformArgPair(matcher.group(2)); + changeCount += addOrUpdate(nestedField, Expressions.truncate(args.getLeft(), args.getRight()), updateSpec, id2field); + break; + } + case "identity": { + changeCount += addOrUpdate(nestedField, Expressions.ref(fieldName), updateSpec, id2field); + break; + } + default: + throw new IllegalStateException("Unexpected value: " + transformer); + } + } + // drop partition + for (PartitionField partitionField : spec.fields()) { + Types.NestedField nestedField = tableSchema.findField(partitionField.sourceId()); + if (nestedField == null || !newPartitionFieldIdSet.contains(nestedField.fieldId())) { + updateSpec.removeField(partitionField.name()); + changeCount++; + } + } + if (changeCount > 0) { + updateSpec.commit(); + return true; + } + return false; + } + + private static int addOrUpdate(Types.NestedField nestedField, Term term, UpdatePartitionSpec updateSpec, + Map id2field) { + PartitionField partitionField = id2field.get(nestedField.fieldId()); + if (partitionField != null) { + if (term instanceof UnboundTransform) { + //noinspection rawtypes + if (((UnboundTransform) term).transform().equals(partitionField.transform())) { + return 0; + } + } else if (term instanceof NamedReference && Transforms.identity().equals(partitionField.transform())) { + return 0; + } + updateSpec.removeField(partitionField.name()); + } + updateSpec.addField(term); + return 1; + } + +} diff --git a/core/src/main/java/kafka/automq/table/utils/TableIdentifierUtil.java b/core/src/main/java/kafka/automq/table/utils/TableIdentifierUtil.java new file mode 100644 index 0000000000..0b0690963f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/utils/TableIdentifierUtil.java @@ -0,0 +1,35 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.utils; + +import org.apache.commons.lang3.StringUtils; +import org.apache.iceberg.catalog.TableIdentifier; + +public class TableIdentifierUtil { + + public static TableIdentifier of(String namespace, String name) { + if (StringUtils.isBlank(namespace)) { + return TableIdentifier.of(name); + } else { + return TableIdentifier.of(namespace, name); + } + } + +} diff --git a/core/src/main/java/kafka/automq/table/worker/BaseDeltaTaskWriter.java b/core/src/main/java/kafka/automq/table/worker/BaseDeltaTaskWriter.java new file mode 100644 index 0000000000..bc870bb248 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/BaseDeltaTaskWriter.java @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.data.InternalRecordWrapper; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.BaseTaskWriter; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.types.TypeUtil; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Set; + +abstract class BaseDeltaTaskWriter extends BaseTaskWriter { + + private final Schema schema; + private final Schema deleteSchema; + private final InternalRecordWrapper wrapper; + private final InternalRecordWrapper keyWrapper; + private final RecordProjection keyProjection; + + BaseDeltaTaskWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema, + Set identifierFieldIds + ) { + super(spec, format, appenderFactory, fileFactory, io, targetFileSize); + this.schema = schema; + this.deleteSchema = TypeUtil.select(schema, new HashSet<>(identifierFieldIds)); + this.wrapper = new InternalRecordWrapper(schema.asStruct()); + this.keyWrapper = new InternalRecordWrapper(deleteSchema.asStruct()); + this.keyProjection = RecordProjection.create(schema, deleteSchema); + } + + abstract RowDataDeltaWriter partition(Record row); + + InternalRecordWrapper wrapper() { + return wrapper; + } + + @Override + public void write(Record row) throws IOException { + Operation op = ((RecordWrapper) row).op(); + RowDataDeltaWriter writer = partition(row); + if (op == Operation.UPDATE || op == Operation.DELETE) { + writer.deleteKey(keyProjection.wrap(row)); + } + if (op == Operation.UPDATE || op == Operation.INSERT) { + writer.write(row); + } + } + + class RowDataDeltaWriter extends BaseEqualityDeltaWriter { + + RowDataDeltaWriter(PartitionKey partition) { + super(partition, schema, deleteSchema); + } + + @Override + protected StructLike asStructLike(Record data) { + return wrapper.wrap(data); + } + + @Override + protected StructLike asStructLikeKey(Record data) { + return keyWrapper.wrap(data); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/EventLoops.java b/core/src/main/java/kafka/automq/table/worker/EventLoops.java new file mode 100644 index 0000000000..cf16f65665 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/EventLoops.java @@ -0,0 +1,214 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.metric.TableTopicMetricsManager; + +import com.automq.stream.s3.metrics.Metrics; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.threads.EventLoop; + +import org.jetbrains.annotations.NotNull; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.PriorityBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import io.netty.util.AbstractReferenceCounted; +import io.netty.util.ReferenceCounted; + +public class EventLoops { + private static final Logger LOGGER = LoggerFactory.getLogger(EventLoops.class); + private final EventLoopWrapper[] eventLoops; + + private final long[] lastRecordNanoTimes; + private final long[] lastTotalBusyTimes; + + public EventLoops(EventLoop[] eventLoops) { + this.eventLoops = new EventLoopWrapper[eventLoops.length]; + for (int i = 0; i < eventLoops.length; i++) { + this.eventLoops[i] = new EventLoopWrapper(eventLoops[i]); + } + this.lastRecordNanoTimes = new long[eventLoops.length]; + long now = System.nanoTime(); + for (int i = 0; i < eventLoops.length; i++) { + lastRecordNanoTimes[i] = now; + } + this.lastTotalBusyTimes = new long[eventLoops.length]; + Threads.COMMON_SCHEDULER.scheduleAtFixedRate(this::logStats, 1, 1, TimeUnit.MINUTES); + } + + public int size() { + return eventLoops.length; + } + + public EventLoopRef leastLoadEventLoop() { + int leastLoad = Integer.MAX_VALUE; + int leastLoadIndex = -1; + for (int i = 0; i < eventLoops.length; i++) { + int load = eventLoops[i].inflight.get(); + if (load < leastLoad) { + leastLoad = load; + leastLoadIndex = i; + } + } + return new EventLoopRef(eventLoops[leastLoadIndex]); + } + + void logStats() { + StringBuilder sb = new StringBuilder(); + long now = System.nanoTime(); + sb.append("EventLoops stats: "); + for (int i = 0; i < eventLoops.length; i++) { + EventLoopWrapper eventLoop = eventLoops[i]; + long totalBusyTime = eventLoop.totalBusyTime; + long lastTotalBusyTime = lastTotalBusyTimes[i]; + long busyTimeDelta = Math.max(totalBusyTime - lastTotalBusyTime, 0); + lastTotalBusyTimes[i] = totalBusyTime; + + long runningTaskStartTime = eventLoop.runningTaskStartTime; + long recordNanoTime = runningTaskStartTime == -1 ? now : runningTaskStartTime; + long lastRecordNanoTime = lastRecordNanoTimes[i]; + lastRecordNanoTimes[i] = recordNanoTime; + + long elapseDelta = Math.max(recordNanoTime - lastRecordNanoTime, 1); + double busyRatio = (double) busyTimeDelta / elapseDelta * 100; + eventLoop.lastBusyRatio = busyRatio; + eventLoop.busyRatioGauge.record(busyRatio); + sb.append(eventLoop.eventLoop.getName()).append(String.format(": %.1f", busyRatio)).append("%, "); + } + LOGGER.info(sb.toString()); + } + + double busyRatio(EventLoop eventLoop) { + for (EventLoopWrapper wrapper : eventLoops) { + if (wrapper.eventLoop == eventLoop) { + return wrapper.lastBusyRatio; + } + } + return 0.0; + } + + public static class EventLoopWrapper { + final EventLoop eventLoop; + + final PriorityBlockingQueue tasks = new PriorityBlockingQueue<>(); + final AtomicInteger inflight = new AtomicInteger(); + volatile long runningTaskStartTime = -1; + volatile long totalBusyTime = 0; + volatile double lastBusyRatio = 0.0; + final Metrics.DoubleGaugeBundle.DoubleGauge busyRatioGauge; + + public EventLoopWrapper(EventLoop eventLoop) { + this.eventLoop = eventLoop; + this.busyRatioGauge = TableTopicMetricsManager.registerEventLoopBusy(eventLoop.getName()); + } + } + + public static class EventLoopRef extends AbstractReferenceCounted implements Executor { + private final EventLoopWrapper eventLoop; + // visible for testing + final AtomicInteger inflight; + + public EventLoopRef(EventLoopWrapper eventLoop) { + this.eventLoop = eventLoop; + this.inflight = eventLoop.inflight; + inflight.incrementAndGet(); + } + + @Override + protected void deallocate() { + inflight.decrementAndGet(); + } + + @Override + public ReferenceCounted touch(Object o) { + return this; + } + + @Override + public void execute(@NotNull Runnable command) { + execute(command, 0); + } + + public CompletableFuture execute(@NotNull Runnable command, long priority) { + CompletableFuture cf = new CompletableFuture<>(); + eventLoop.tasks.add(new PriorityTask(() -> { + eventLoop.runningTaskStartTime = System.nanoTime(); + try { + command.run(); + } finally { + //noinspection NonAtomicOperationOnVolatileField + eventLoop.totalBusyTime += System.nanoTime() - eventLoop.runningTaskStartTime; + eventLoop.runningTaskStartTime = -1L; + } + }, (int) priority, cf)); + eventLoop.eventLoop.execute(() -> { + CompletableFuture headTaskCf = null; + try { + PriorityTask headTask = eventLoop.tasks.take(); + headTaskCf = headTask.cf; + headTask.task.run(); + headTask.cf.complete(null); + } catch (Throwable e) { + if (headTaskCf != null) { + headTaskCf.completeExceptionally(e); + } + throw new RuntimeException(e); + } + }); + return cf; + } + } + + static final class PriorityTask implements Comparable { + private final Runnable task; + private final int priority; + private final CompletableFuture cf; + + PriorityTask(Runnable task, int priority, CompletableFuture cf) { + this.task = task; + this.priority = priority; + this.cf = cf; + } + + public Runnable task() { + return task; + } + + public int priority() { + return priority; + } + + public CompletableFuture cf() { + return cf; + } + + @Override + public int compareTo(@NotNull PriorityTask o) { + return Integer.compare(priority, o.priority); + } + } + +} diff --git a/core/src/main/java/kafka/automq/table/worker/FlushMode.java b/core/src/main/java/kafka/automq/table/worker/FlushMode.java new file mode 100644 index 0000000000..1e19b5d89f --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/FlushMode.java @@ -0,0 +1,25 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +public enum FlushMode { + FLUSH, + COMPLETE +} diff --git a/core/src/main/java/kafka/automq/table/worker/IcebergTableManager.java b/core/src/main/java/kafka/automq/table/worker/IcebergTableManager.java new file mode 100644 index 0000000000..5d5775d434 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/IcebergTableManager.java @@ -0,0 +1,358 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.utils.PartitionUtil; + +import com.google.common.annotations.VisibleForTesting; + +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.UpdateSchema; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.SupportsNamespaces; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.exceptions.AlreadyExistsException; +import org.apache.iceberg.exceptions.NoSuchNamespaceException; +import org.apache.iceberg.exceptions.NoSuchTableException; +import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.Tasks; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +public class IcebergTableManager { + private static final Logger LOGGER = LoggerFactory.getLogger(IcebergTableManager.class); + private final Catalog catalog; + private final TableIdentifier tableId; + private final WorkerConfig config; + private volatile Table table; + private volatile PartitionSpec partitionSpec; + + public IcebergTableManager(Catalog catalog, TableIdentifier tableId, WorkerConfig config) { + this.catalog = catalog; + this.tableId = tableId; + this.config = config; + } + + public Table getTableOrCreate(Schema schema) { + Table currentTable = table; + if (currentTable == null) { + synchronized (this) { + currentTable = table; + if (currentTable == null) { + table = currentTable = getTableOrCreate(schema, 1); + partitionSpec = currentTable.spec(); + } + } + } + return currentTable; + } + + public Table getTableOrCreate(Schema schema, int retries) { + AtomicReference result = new AtomicReference<>(); + Tasks.range(1) + .retry(retries) + .run(notUsed -> { + try { + result.set(catalog.loadTable(tableId)); + } catch (NoSuchTableException e) { + if (catalog instanceof SupportsNamespaces) { + SupportsNamespaces namespaceCatalog = (SupportsNamespaces) catalog; + try { + namespaceCatalog.loadNamespaceMetadata(tableId.namespace()); + } catch (NoSuchNamespaceException e2) { + LOGGER.info("Namespace {} does not exist, creating", tableId.namespace()); + try { + ((SupportsNamespaces) catalog).createNamespace(tableId.namespace()); + } catch (AlreadyExistsException e3) { + LOGGER.info("Namespace {} already exists", tableId.namespace()); + } + } + } + try { + PartitionSpec spec = PartitionUtil.buildPartitionSpec(config.partitionBy(), schema); + Map options = new HashMap<>(); + options.put(TableProperties.METADATA_DELETE_AFTER_COMMIT_ENABLED, "true"); + options.put(TableProperties.OBJECT_STORE_ENABLED, "true"); + LOGGER.info("Table {} does not exist, create with schema={}, partition={}", tableId, schema, spec); + result.set(catalog.createTable(tableId, schema, spec, options)); + } catch (AlreadyExistsException e1) { + LOGGER.info("Table {} already exists", tableId); + result.set(catalog.loadTable(tableId)); + } + } + }); + return result.get(); + } + + public boolean handleSchemaChangesWithFlush(Schema schema, FlushAction flush) throws IOException { + Table currentTable = getTableOrCreate(schema); + List changes = checkSchemaChanges(currentTable, schema); + if (changes.isEmpty()) { + return false; + } + + flush.perform(); + + applySchemaChange(currentTable, changes); + return true; + } + + public TableIdentifier tableId() { + return tableId; + } + + public PartitionSpec spec() { + return partitionSpec; + } + + public void reset() { + table = null; + } + + /** + * Check schema changes between the current record and the table schema + */ + @VisibleForTesting + protected List checkSchemaChanges(Table table, Schema currentSchema) { + return collectSchemaChanges(currentSchema, table); + } + + /** + * Apply schema changes to the table + * + * @param changes list of schema changes + */ + @VisibleForTesting + protected synchronized void applySchemaChange(Table table, List changes) { + LOGGER.info("Applying schema changes to table {}, changes {}", tableId, changes.stream().map(c -> c.getType() + ":" + c.getColumnFullName()).toList()); + Tasks.range(1) + .retry(2) + .run(notUsed -> applyChanges(table, changes)); + table.refresh(); + } + + private static UpdateSchema applySchemaChange(UpdateSchema updateSchema, SchemaChange change) { + switch (change.getType()) { + case ADD_COLUMN: + if (change.getParentName() == null) { + return updateSchema + .addColumn(change.getColumnFullName(), change.getNewType()); + } else { + return updateSchema + .addColumn(change.getParentName(), change.getColumnName(), change.getNewType()); + } + case MAKE_OPTIONAL: + return updateSchema + .makeColumnOptional(change.getColumnFullName()); + case PROMOTE_TYPE: + return updateSchema + .updateColumn(change.getColumnFullName(), change.getNewType().asPrimitiveType()); + default: + return updateSchema; + } + } + + private static boolean shouldSkipChange(org.apache.iceberg.Schema schema, SchemaChange change) { + Types.NestedField field = schema.findField(change.getColumnFullName()); + switch (change.getType()) { + case ADD_COLUMN: { + return field != null; + } + case MAKE_OPTIONAL: { + return field != null && field.isOptional(); + } + case PROMOTE_TYPE: { + return field != null && field.type().equals(change.getNewType()); + } + default: { + return false; + } + } + } + + protected List collectSchemaChanges(Schema currentSchema, Table table) { + Schema tableSchema = table.schema(); + List changes = new ArrayList<>(); + + for (Types.NestedField currentField : currentSchema.columns()) { + collectFieldChanges(currentField, null, tableSchema, changes); + } + + for (Types.NestedField tableField : tableSchema.columns()) { + collectRemovedField(tableField, null, currentSchema, changes); + } + return changes; + } + + private void collectRemovedField(Types.NestedField tableField, String parentName, Schema currentSchema, + List changes) { + String fieldName = tableField.name(); + String fullFieldName = parentName == null ? fieldName : parentName + "." + fieldName; + Types.NestedField currentField = currentSchema.findField(fullFieldName); + + // if field doesn't exist in current schema and it's not a struct, mark it as optional (soft removal) + if (currentField == null && !tableField.isOptional()) { + changes.add(new SchemaChange(SchemaChange.ChangeType.MAKE_OPTIONAL, fieldName, + null, parentName)); + return; + } + // if it is a nested field, recursively process subfields + if (tableField.type().isStructType()) { + collectRemovedStructFields(tableField.type().asStructType().fields(), fullFieldName, currentSchema, changes); + } else if (isStructList(tableField.type())) { + collectRemovedStructFields(tableField.type().asListType().elementType().asStructType().fields(), + fullFieldName + ".element", currentSchema, changes); + } else if (isStructMap(tableField.type())) { + collectRemovedStructFields(tableField.type().asMapType().valueType().asStructType().fields(), + fullFieldName + ".value", currentSchema, changes); + } + } + + private void collectFieldChanges(Types.NestedField currentField, String parentName, Schema tableSchema, + List changes) { + String fieldName = currentField.name(); + String fullFieldName = parentName == null ? fieldName : parentName + "." + fieldName; + Types.NestedField tableField = tableSchema.findField(fullFieldName); + + if (tableField == null) { + changes.add(new SchemaChange(SchemaChange.ChangeType.ADD_COLUMN, fieldName, + currentField.type(), parentName)); + return; + } else { + Type currentType = currentField.type(); + Type tableType = tableField.type(); + if (currentType.isStructType() && tableType.isStructType()) { + collectStructFieldChanges(currentType.asStructType().fields(), fullFieldName, tableSchema, changes); + collectOptionalFieldChanges(currentField, parentName, changes, tableField, fieldName); + } else if (isStructList(currentType) && isStructList(tableType)) { + collectStructFieldChanges(currentType.asListType().elementType().asStructType().fields(), + fullFieldName + ".element", tableSchema, changes); + } else if (isStructMap(currentType) && isStructMap(tableType)) { + collectStructFieldChanges(currentType.asMapType().valueType().asStructType().fields(), + fullFieldName + ".value", tableSchema, changes); + } else if (!currentType.isStructType() && !tableType.isStructType()) { + collectOptionalFieldChanges(currentField, parentName, changes, tableField, fieldName); + + if (!tableType.equals(currentType) && canPromoteType(tableType, currentType)) { + changes.add(new SchemaChange(SchemaChange.ChangeType.PROMOTE_TYPE, fieldName, currentType, parentName)); + } + } + } + } + + private static void collectOptionalFieldChanges(Types.NestedField currentField, String parentName, List changes, Types.NestedField tableField, String fieldName) { + if (!tableField.isOptional() && currentField.isOptional()) { + changes.add(new SchemaChange(SchemaChange.ChangeType.MAKE_OPTIONAL, fieldName, null, parentName)); + } + } + + private void collectStructFieldChanges(List currentSubFields, String parentFullName, + Schema tableSchema, List changes) { + for (Types.NestedField currentSubField : currentSubFields) { + collectFieldChanges(currentSubField, parentFullName, tableSchema, changes); + } + } + + private void collectRemovedStructFields(List tableSubFields, String parentFullName, + Schema currentSchema, List changes) { + for (Types.NestedField tableSubField : tableSubFields) { + collectRemovedField(tableSubField, parentFullName, currentSchema, changes); + } + } + + private boolean isStructList(Type type) { + return type.typeId() == Type.TypeID.LIST && type.asListType().elementType().isStructType(); + } + + private boolean isStructMap(Type type) { + return type.typeId() == Type.TypeID.MAP && type.asMapType().valueType().isStructType(); + } + + private boolean canPromoteType(Type oldType, Type newType) { + if (oldType.typeId() == Type.TypeID.INTEGER && newType.typeId() == Type.TypeID.LONG) { + return true; + } + return oldType.typeId() == Type.TypeID.FLOAT && newType.typeId() == Type.TypeID.DOUBLE; + } + + private void applyChanges(Table table, List changes) { + table.refresh(); + UpdateSchema updateSchema = table.updateSchema(); + changes.stream().filter(c -> !shouldSkipChange(table.schema(), c)) + .forEach(c -> applySchemaChange(updateSchema, c)); + updateSchema.commit(); + } + + @FunctionalInterface + public interface FlushAction { + void perform() throws IOException; + } + + static class SchemaChange { + private final ChangeType type; + private final String columnName; + private final Type newType; + private final String parentName; // For nested fields + + public SchemaChange(ChangeType type, String columnName, Type newType, String parentName) { + this.type = type; + this.columnName = columnName; + this.newType = newType; + this.parentName = parentName; + } + + public ChangeType getType() { + return type; + } + + public String getColumnName() { + return columnName; + } + + public Type getNewType() { + return newType; + } + + public String getParentName() { + return parentName; + } + + public String getColumnFullName() { + return parentName == null ? columnName : parentName + "." + columnName; + } + + enum ChangeType { + ADD_COLUMN, + MAKE_OPTIONAL, + PROMOTE_TYPE + } + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/IcebergWriter.java b/core/src/main/java/kafka/automq/table/worker/IcebergWriter.java new file mode 100644 index 0000000000..448d1d1199 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/IcebergWriter.java @@ -0,0 +1,532 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.binder.RecordBinder; +import kafka.automq.table.events.PartitionMetric; +import kafka.automq.table.events.TopicMetric; +import kafka.automq.table.process.DataError; +import kafka.automq.table.process.ProcessingResult; +import kafka.automq.table.process.RecordProcessor; +import kafka.automq.table.process.exception.RecordProcessorException; + +import org.apache.kafka.server.record.ErrorsTolerance; + +import com.automq.stream.s3.metrics.TimerUtil; +import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; +import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; +import com.automq.stream.s3.network.NetworkBandwidthLimiter; +import com.automq.stream.s3.network.ThrottleStrategy; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.LogSuppressor; + +import org.apache.avro.generic.GenericRecord; +import org.apache.commons.lang3.StringUtils; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.io.UnpartitionedWriter; +import org.apache.iceberg.io.WriteResult; +import org.apache.iceberg.types.TypeUtil; +import org.apache.iceberg.types.Types; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import static java.util.stream.Collectors.toSet; + +public class IcebergWriter implements Writer { + private static final Logger LOGGER = LoggerFactory.getLogger(IcebergWriter.class); + private static final int TARGET_FILE_SIZE = 64 * 1024 * 1024; + private static final LogSuppressor INVALID_DATA_LOGGER = new LogSuppressor(LOGGER, 60000); + private static final int WRITER_RESULT_SIZE_LIMIT = 5; + final List results = new ArrayList<>(); + private final TableIdentifier tableId; + private final RecordProcessor processor; + private final IcebergTableManager icebergTableManager; + private final Map offsetRangeMap = new HashMap<>(); + private TaskWriter writer; + private int recordCount = 0; + private long dirtyBytes = 0; + private long fieldCount = 0; + private final Map metrics = new HashMap<>(); + private Status status = Status.WRITABLE; + private final WorkerConfig config; + private final boolean deltaWrite; + private RecordBinder binder; + private String lastSchemaIdentity; + private final NetworkBandwidthLimiter outboundLimiter; + + + public IcebergWriter(IcebergTableManager icebergTableManager, RecordProcessor processor, WorkerConfig config) { + this.tableId = icebergTableManager.tableId(); + this.icebergTableManager = icebergTableManager; + this.processor = processor; + this.config = config; + this.deltaWrite = StringUtils.isNoneBlank(config.cdcField()) || config.upsertEnable(); + this.outboundLimiter = GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.OUTBOUND); + } + + @Override + public void write(int partition, org.apache.kafka.common.record.Record kafkaRecord) throws IOException { + if (status != Status.WRITABLE) { + throw new IOException(String.format("The writer %s isn't in WRITABLE status, current status is %s", this, status)); + } + try { + if (write0(partition, kafkaRecord)) { + recordCount++; + dirtyBytes += kafkaRecord.sizeInBytes(); + } + offsetRangeMap.compute(partition, (k, v) -> { + if (v == null) { + throw new IllegalArgumentException(String.format("The partition %s initial offset is not set", partition)); + } + long recordOffset = kafkaRecord.offset(); + if (recordOffset < v.end) { + throw new IllegalArgumentException(String.format("The record offset[%s] is less than the end[%s]", recordOffset, v.end)); + } + v.end = recordOffset + 1; + return v; + }); + } catch (Throwable e) { + LOGGER.error("[WRITE_FAIL],{}", this, e); + status = Status.ERROR; + throw new IOException(e); + } + } + + @Override + public String toString() { + return tableId.toString(); + } + + protected boolean write0(int partition, + org.apache.kafka.common.record.Record kafkaRecord) throws IOException, RecordProcessorException { + ProcessingResult result = processor.process(partition, kafkaRecord); + + if (!result.isSuccess()) { + DataError error = result.getError(); + String recordContext = buildRecordContext(partition, kafkaRecord); + String errorMsg = String.format("Data processing failed for record: %s", recordContext); + + if (config.errorsTolerance() == ErrorsTolerance.ALL + || (DataError.ErrorType.DATA_ERROR.equals(error.getType()) + && config.errorsTolerance().equals(ErrorsTolerance.INVALID_DATA))) { + INVALID_DATA_LOGGER.warn("[INVALID_DATA],{}", this, error.getCause()); + return false; + } else { + throw new RecordProcessorException(errorMsg + " - " + error.getDetailedMessage(), error.getCause()); + } + } + + GenericRecord finalRecord = result.getFinalRecord(); + + RecordBinder currentBinder = this.binder; + // first write + if (currentBinder == null) { + currentBinder = new RecordBinder(finalRecord); + } + + // schema change + if (!result.getFinalSchemaIdentity().equals(lastSchemaIdentity)) { + Schema icebergSchema = new RecordBinder(finalRecord).getIcebergSchema(); + // compare table schema and evolution + icebergTableManager.handleSchemaChangesWithFlush( + icebergSchema, + this::flush + ); + + // update Binder + Table table = icebergTableManager.getTableOrCreate(icebergSchema); + currentBinder = currentBinder.createBinderForNewSchema(table.schema(), finalRecord.getSchema()); + lastSchemaIdentity = result.getFinalSchemaIdentity(); + } + Record record = currentBinder.bind(finalRecord); + this.binder = currentBinder; + + recordMetric(partition, kafkaRecord.timestamp()); + + waitForNetworkPermit(1); + + TaskWriter writer = getWriter(record.struct()); + + + if (deltaWrite) { + writer.write(new RecordWrapper(record, config.cdcField(), config.upsertEnable())); + } else { + writer.write(record); + } + return true; + } + + @Override + public CompletableFuture flush(FlushMode flushMode, ExecutorService flushExecutor, Executor eventLoop) { + switch (status) { + case WRITABLE: + status = Status.FLUSHING; + break; + case FLUSHING: + return FutureUtil.failedFuture(new IOException("The writer is already flushing")); + case COMPLETED: + return CompletableFuture.completedFuture(null); + case ERROR: + return FutureUtil.failedFuture(errorStateException()); + default: + return FutureUtil.failedFuture(errorStateException()); + } + CompletableFuture cf = new CompletableFuture<>(); + flushExecutor.submit(() -> { + try { + TimerUtil timerUtil = new TimerUtil(); + boolean hasData = flush(); + if (flushMode == FlushMode.COMPLETE) { + complete(); + } + if (hasData && LOGGER.isTraceEnabled()) { + LOGGER.trace("[TABLE_FLUSH],{},{}ms", this, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); + } + cf.complete(null); + } catch (Throwable e) { + LOGGER.error("[DATA_FILE_FLUSH_FAIL],{}", this, e); + cf.completeExceptionally(e); + } + }); + return cf.thenAcceptAsync(nil -> { + if (status == Status.FLUSHING) { + // the status may become ERROR or COMPLETED + status = Status.WRITABLE; + } + }, eventLoop); + } + + @Override + public void abort() throws IOException { + } + + @Override + public List complete() throws IOException { + check(); + if (status == Status.COMPLETED) { + return results; + } + flush(); + status = Status.COMPLETED; + return results; + } + + @Override + public List results() { + if (status != Status.COMPLETED) { + throw new IllegalStateException("The writer isn't completed, current status is " + status); + } + return results; + } + + @Override + public boolean isCompleted() { + return status == Status.COMPLETED; + } + + @Override + public boolean isFull() { + return results.size() >= WRITER_RESULT_SIZE_LIMIT; + } + + @Override + public Map getOffsets() { + return offsetRangeMap; + } + + @Override + public OffsetRange getOffset(int partition) { + return offsetRangeMap.get(partition); + } + + @Override + public void setOffset(int partition, long offset) { + if (offsetRangeMap.containsKey(partition)) { + throw new IllegalArgumentException(String.format("The partition %s initial offset is already set", partition)); + } + offsetRangeMap.put(partition, new OffsetRange(offset)); + metrics.put(partition, new Metric()); + } + + @Override + public void setEndOffset(int partition, long offset) { + offsetRangeMap.compute(partition, (k, v) -> { + if (v == null) { + throw new IllegalArgumentException(String.format("The partition %s initial offset is not set", partition)); + } + v.end = offset; + return v; + }); + } + + @Override + public long dirtyBytes() { + return dirtyBytes; + } + + public void updateWatermark(int partition, long watermark) { + recordMetric(partition, watermark); + } + + @Override + public TopicMetric topicMetric() { + return new TopicMetric(fieldCount); + } + + @Override + public Map partitionMetrics() { + return metrics.entrySet().stream().collect(Collectors.toMap( + Map.Entry::getKey, + e -> new PartitionMetric(e.getKey(), e.getValue().watermark) + )); + } + + @Override + public int targetFileSize() { + return TARGET_FILE_SIZE; + } + + private void check() throws IOException { + if (status == Status.ERROR) { + throw errorStateException(); + } + } + + private IOException errorStateException() { + return new IOException(String.format("The writer %s is in error state", this)); + } + + private TaskWriter getWriter(Types.StructType prototype) { + if (writer == null) { + writer = writer(icebergTableManager.getTableOrCreate(prototype.asSchema())); + } + return writer; + } + + private TaskWriter writer(Table table) { + FileAppenderFactory appenderFactory; + + Set identifierFieldIds = table.schema().identifierFieldIds(); + if (!config.idColumns().isEmpty()) { + identifierFieldIds = config.idColumns().stream() + .map(colName -> table.schema().findField(colName).fieldId()) + .collect(toSet()); + } + // Use a consistent partition spec instead of retrieve from table in real times. + PartitionSpec spec = icebergTableManager.spec(); + if (identifierFieldIds.isEmpty()) { + appenderFactory = + new GenericAppenderFactory(table.schema(), spec, null, null, null) + .setAll(table.properties()); + } else { + appenderFactory = + new GenericAppenderFactory( + table.schema(), + spec, + identifierFieldIds.stream().mapToInt(i -> i).toArray(), + TypeUtil.select(table.schema(), identifierFieldIds), + null) + .setAll(table.properties()); + } + + // (partition ID + task ID + operation ID) must be unique + OutputFileFactory fileFactory = + OutputFileFactory.builderFor(table, 1, System.currentTimeMillis()) + .defaultSpec(spec) + .operationId(UUID.randomUUID().toString()) + .format(FileFormat.PARQUET) + .build(); + + TaskWriter writer; + if (spec.isUnpartitioned()) { + if (!deltaWrite) { + writer = + new UnpartitionedWriter<>( + spec, FileFormat.PARQUET, appenderFactory, fileFactory, table.io(), TARGET_FILE_SIZE); + } else { + writer = + new UnpartitionedDeltaWriter( + spec, + FileFormat.PARQUET, + appenderFactory, + fileFactory, + table.io(), + TARGET_FILE_SIZE, + table.schema(), + identifierFieldIds); + } + } else { + if (!deltaWrite) { + writer = + new PartitionedWriter( + table.spec(), + FileFormat.PARQUET, + appenderFactory, + fileFactory, + table.io(), + TARGET_FILE_SIZE, + table.schema()); + } else { + writer = + new PartitionedDeltaWriter( + table.spec(), + FileFormat.PARQUET, + appenderFactory, + fileFactory, + table.io(), + TARGET_FILE_SIZE, + table.schema(), + identifierFieldIds); + } + } + return writer; + } + + private boolean flush() throws IOException { + if (status == Status.ERROR) { + throw errorStateException(); + } + try { + if (writer == null) { + return false; + } + if (recordCount == 0) { + return false; + } + // Complete writer first, then collect statistics only if successful + WriteResult writeResult = this.writer.complete(); + results.add(writeResult); + recordNetworkCost(writeResult); + + // Collect field count statistics from the binder after successful completion + if (binder != null) { + fieldCount += binder.getAndResetFieldCount(); + } + + this.writer = null; + recordCount = 0; + dirtyBytes = 0; + return true; + } catch (Throwable e) { + status = Status.ERROR; + throw e; + } + } + + private void recordNetworkCost(WriteResult writeResult) { + final long totalBytes = calculateWriteResultBytes(writeResult); + if (totalBytes <= 0) { + LOGGER.warn("[NETWORK_LIMITER_RECORD_INVALID_BYTES],{},bytes={}", this, totalBytes); + return; + } + try { + waitForNetworkPermit(totalBytes); + } catch (IOException e) { + LOGGER.warn("[NETWORK_LIMITER_RECORD_FAIL],{},bytes={}", this, totalBytes, e); + } + } + + private long calculateWriteResultBytes(WriteResult writeResult) { + long bytes = 0L; + for (DataFile file : writeResult.dataFiles()) { + bytes += Math.max(file.fileSizeInBytes(), 0L); + } + for (DeleteFile file : writeResult.deleteFiles()) { + bytes += Math.max(file.fileSizeInBytes(), 0L); + } + return bytes; + } + + private void waitForNetworkPermit(long size) throws IOException { + try { + outboundLimiter.consumeBlocking(ThrottleStrategy.ICEBERG_WRITE, size); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn("[NETWORK_LIMITER_PERMIT_FAIL],{}", this, e); + throw new IOException("Failed to acquire outbound network permit", e); + } catch (ExecutionException e) { + Throwable cause = e.getCause() == null ? e : e.getCause(); + LOGGER.warn("[NETWORK_LIMITER_PERMIT_ERROR],{}", this, cause); + throw new IOException("Failed to acquire outbound network permit", cause); + } catch (CancellationException e) { + LOGGER.warn("[NETWORK_LIMITER_PERMIT_ERROR],{}", this, e); + throw new IOException("Failed to acquire outbound network permit", e); + } catch (RuntimeException e) { + LOGGER.warn("[NETWORK_LIMITER_PERMIT_FAIL],{}", this, e); + throw new IOException("Failed to acquire outbound network permit", e); + } + } + + private void recordMetric(int partition, long timestamp) { + Metric metric = metrics.get(partition); + if (metric.watermark < timestamp) { + metric.watermark = timestamp; + } + } + + /** + * Builds a descriptive context string for a Kafka record to include in error messages. + */ + private String buildRecordContext(int partition, org.apache.kafka.common.record.Record kafkaRecord) { + return String.format("topic=%s, partition=%d, offset=%d, timestamp=%d", + tableId.name(), + partition, + kafkaRecord.offset(), + kafkaRecord.timestamp()); + } + + static class Metric { + long watermark = -1L; + } + + enum Status { + WRITABLE, + FLUSHING, + COMPLETED, + ERROR, + } + +} diff --git a/core/src/main/java/kafka/automq/table/worker/IcebergWriterFactory.java b/core/src/main/java/kafka/automq/table/worker/IcebergWriterFactory.java new file mode 100644 index 0000000000..a9877db7ea --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/IcebergWriterFactory.java @@ -0,0 +1,58 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.process.RecordProcessorFactory; + +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; + +public class IcebergWriterFactory implements WriterFactory { + private final TableIdentifier tableIdentifier; + private final IcebergTableManager icebergTableManager; + private final RecordProcessorFactory recordProcessorFactory; + private final WorkerConfig config; + private final String topic; + + public IcebergWriterFactory(Catalog catalog, TableIdentifier tableIdentifier, RecordProcessorFactory recordProcessorFactory, WorkerConfig config, String topic) { + this.topic = topic; + this.tableIdentifier = tableIdentifier; + this.icebergTableManager = new IcebergTableManager(catalog, tableIdentifier, config); + this.recordProcessorFactory = recordProcessorFactory; + this.config = config; + } + + @Override + public Writer newWriter() { + return new IcebergWriter(icebergTableManager, + recordProcessorFactory.create(config, topic), config); + } + + @Override + public PartitionSpec partitionSpec() { + return icebergTableManager.spec(); + } + + @Override + public void reset() { + icebergTableManager.reset(); + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/Operation.java b/core/src/main/java/kafka/automq/table/worker/Operation.java new file mode 100644 index 0000000000..882641b7bc --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/Operation.java @@ -0,0 +1,25 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +public enum Operation { + INSERT, + UPDATE, + DELETE +} diff --git a/core/src/main/java/kafka/automq/table/worker/PartitionWriteTask.java b/core/src/main/java/kafka/automq/table/worker/PartitionWriteTask.java new file mode 100644 index 0000000000..c14222f551 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/PartitionWriteTask.java @@ -0,0 +1,174 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.cluster.Partition; +import kafka.log.streamaspect.ReadHint; + +import org.apache.kafka.common.record.PooledResource; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.utils.BufferSupplier; +import org.apache.kafka.storage.internals.log.FetchDataInfo; +import org.apache.kafka.storage.internals.log.FetchIsolation; + +import com.automq.stream.s3.metrics.TimerUtil; +import com.automq.stream.utils.AsyncSemaphore; +import com.automq.stream.utils.Systems; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +@SuppressWarnings({"CyclomaticComplexity", "NPathComplexity"}) +class PartitionWriteTask { + private static final Logger LOGGER = LoggerFactory.getLogger(PartitionWriteTask.class); + private static final AsyncSemaphore READ_LIMITER = new AsyncSemaphore(Systems.CPU_CORES * 100L * 1024 * 1024); + private final long startTimestamp = System.currentTimeMillis(); + final Partition partition; + final long taskStartOffset; + final long endOffset; + final PartitionWriteTaskContext ctx; + final CompletableFuture cf = new CompletableFuture<>(); + + public PartitionWriteTask(Partition partition, long startOffset, long endOffset, PartitionWriteTaskContext ctx) { + this.partition = partition; + this.taskStartOffset = Math.max(startOffset, partition.log().get().logStartOffset()); + this.endOffset = endOffset; + this.ctx = ctx; + } + + public CompletableFuture run() { + run0(taskStartOffset); + return cf; + } + + private void run0(long startOffset) { + int readSize = 1024 * 1024; + // limit the read direct memory usage. + READ_LIMITER.acquire(readSize, () -> { + CompletableFuture readReleaseCf = new CompletableFuture<>(); + readAsync(startOffset, readSize) + .thenCompose(rst -> ctx.eventLoop.execute(() -> handleReadResult(startOffset, rst), ctx.priority)) + .exceptionally(ex -> { + ctx.requireReset = true; + LOGGER.error("Error in read task {}", this, ex); + cf.complete(null); + return null; + }).whenComplete((nil, ex) -> readReleaseCf.complete(null)); + return readReleaseCf; + }, ctx.eventLoop); + } + + private CompletableFuture readAsync(long startOffset, int readSize) { + try { + ReadHint.markReadAll(); + // FIXME: wrap readAsync the exception in the future. + return partition.log().get() + .readAsync(startOffset, readSize, FetchIsolation.TXN_COMMITTED, true); + } catch (Throwable ex) { + // When the partition is closed, the readAsync will throw exception. + return CompletableFuture.failedFuture(ex); + } + } + + private void handleReadResult(long startOffset, FetchDataInfo rst) { + BufferSupplier bufferSupplier = BufferSupplier.create(); + try { + if (ctx.requireReset) { + cf.complete(null); + return; + } + if (!ctx.lastFlushCf.isDone()) { + // Avoid concurrent write and flush to the same writer + ctx.lastFlushCf.whenComplete((nil, ex) -> { + if (ex != null) { + ctx.requireReset = true; + cf.complete(null); + } else { + ctx.eventLoop.execute(() -> handleReadResult(startOffset, rst), ctx.priority); + } + }); + return; + } + TimerUtil timer = new TimerUtil(); + + long writeSize = 0; + RecordsIterator it = new RecordsIterator(startOffset, rst, bufferSupplier); + long nextOffset = startOffset; + while (it.hasNext()) { + Record record = it.next(); + long recordOffset = record.offset(); + if (recordOffset >= startOffset && recordOffset < endOffset) { + ctx.writer.write(partition.partitionId(), record); + writeSize += record.sizeInBytes(); + nextOffset = recordOffset + 1; + } else if (recordOffset >= endOffset) { + nextOffset = endOffset; + // Abort transactions might occupy the offsets + ctx.writer.setEndOffset(partition.partitionId(), nextOffset); + break; + } + } + if (!it.hasNext()) { + nextOffset = Math.min(endOffset, it.nextOffset()); + // Abort transactions might occupy the offsets + ctx.writer.setEndOffset(partition.partitionId(), nextOffset); + } + + if (ctx.writer.dirtyBytes() >= ctx.writer.targetFileSize()) { + ctx.lastFlushCf = ctx.writer.flush(FlushMode.FLUSH, ctx.flushExecutors, ctx.eventLoop); + } + if (nextOffset != startOffset && nextOffset < endOffset) { + // launch next round read until read to the end. + long finalNextOffset = nextOffset; + ctx.eventLoop.execute(() -> run0(finalNextOffset)); + } else { + if (nextOffset == partition.log().get().highWatermark()) { + // Update the timestamp to the start timestamp to avoid that situation low traffic topic record old timestamps. + ctx.writer.updateWatermark(partition.partitionId(), startTimestamp); + } + ctx.lastFlushCf.whenCompleteAsync((nil, ex) -> cf.complete(null), ctx.eventLoop); + } + if (nextOffset != startOffset) { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[HANDLE_READ_RESULT],{},{}-{},{}ms", this, startOffset, nextOffset, timer.elapsedAs(TimeUnit.MILLISECONDS)); + } + ctx.recordWriteSize(writeSize); + } + } catch (Throwable e) { + ctx.requireReset = true; + LOGGER.error("[HANDLE_READ_RESULT_FAIL],{},{}-{}", this, startOffset, endOffset, e); + cf.complete(null); + } finally { + if (rst.records instanceof PooledResource) { + ((PooledResource) rst.records).release(); + } + bufferSupplier.close(); + } + } + + @Override + public String toString() { + return partition.topic() + "-" + partition.partitionId(); + } + +} diff --git a/core/src/main/java/kafka/automq/table/worker/PartitionWriteTaskContext.java b/core/src/main/java/kafka/automq/table/worker/PartitionWriteTaskContext.java new file mode 100644 index 0000000000..761a41813e --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/PartitionWriteTaskContext.java @@ -0,0 +1,50 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; + +class PartitionWriteTaskContext { + boolean requireReset; + CompletableFuture lastFlushCf; + + final Writer writer; + final EventLoops.EventLoopRef eventLoop; + final ExecutorService flushExecutors; + final WorkerConfig config; + final long priority; + long writeSize = 0; + + public PartitionWriteTaskContext(Writer writer, EventLoops.EventLoopRef eventLoop, ExecutorService flushExecutors, WorkerConfig config, long priority) { + this.requireReset = false; + this.lastFlushCf = CompletableFuture.completedFuture(null); + + this.writer = writer; + this.eventLoop = eventLoop; + this.flushExecutors = flushExecutors; + this.config = config; + this.priority = priority; + } + + public void recordWriteSize(long size) { + writeSize += size; + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/PartitionedDeltaWriter.java b/core/src/main/java/kafka/automq/table/worker/PartitionedDeltaWriter.java new file mode 100644 index 0000000000..da25831d46 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/PartitionedDeltaWriter.java @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.util.Tasks; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +public class PartitionedDeltaWriter extends BaseDeltaTaskWriter { + private final PartitionKey partitionKey; + + private final Map writers = new HashMap<>(); + + PartitionedDeltaWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema, + Set identifierFieldIds + ) { + super( + spec, + format, + appenderFactory, + fileFactory, + io, + targetFileSize, + schema, + identifierFieldIds + ); + this.partitionKey = new PartitionKey(spec, schema); + } + + @Override + RowDataDeltaWriter partition(Record row) { + partitionKey.partition(wrapper().wrap(row)); + + RowDataDeltaWriter writer = writers.get(partitionKey); + if (writer == null) { + // NOTICE: we need to copy a new partition key here, in case of messing up the keys in + // writers. + PartitionKey copiedKey = partitionKey.copy(); + writer = new RowDataDeltaWriter(copiedKey); + writers.put(copiedKey, writer); + } + + return writer; + } + + @Override + public void close() { + try { + Tasks.foreach(writers.values()) + .throwFailureWhenFinished() + .noRetry() + .run(RowDataDeltaWriter::close, IOException.class); + + writers.clear(); + } catch (IOException e) { + throw new UncheckedIOException("Failed to close equality delta writer", e); + } + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/PartitionedWriter.java b/core/src/main/java/kafka/automq/table/worker/PartitionedWriter.java new file mode 100644 index 0000000000..91993624c3 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/PartitionedWriter.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionKey; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.InternalRecordWrapper; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.PartitionedFanoutWriter; + +class PartitionedWriter extends PartitionedFanoutWriter { + + private final PartitionKey partitionKey; + private final InternalRecordWrapper wrapper; + + PartitionedWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema) { + super(spec, format, appenderFactory, fileFactory, io, targetFileSize); + this.partitionKey = new PartitionKey(spec, schema); + this.wrapper = new InternalRecordWrapper(schema.asStruct()); + } + + @Override + protected PartitionKey partition(Record row) { + partitionKey.partition(wrapper.wrap(row)); + return partitionKey; + } +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/table/worker/RecordProjection.java b/core/src/main/java/kafka/automq/table/worker/RecordProjection.java new file mode 100644 index 0000000000..c3e18009be --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/RecordProjection.java @@ -0,0 +1,202 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +import com.google.common.base.Preconditions; + +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.types.Types.ListType; +import org.apache.iceberg.types.Types.MapType; +import org.apache.iceberg.types.Types.NestedField; +import org.apache.iceberg.types.Types.StructType; + +import java.util.List; +import java.util.Map; + +/** + * This is modified from {@link org.apache.iceberg.util.StructProjection} to support record types. + */ +public class RecordProjection implements Record { + + /** + * Creates a projecting wrapper for {@link Record} rows. + * + *

This projection does not work with repeated types like lists and maps. + * + * @param dataSchema schema of rows wrapped by this projection + * @param projectedSchema result schema of the projected rows + * @return a wrapper to project rows + */ + public static RecordProjection create(Schema dataSchema, Schema projectedSchema) { + return new RecordProjection(dataSchema.asStruct(), projectedSchema.asStruct()); + } + + private final StructType type; + private final int[] positionMap; + private final RecordProjection[] nestedProjections; + private Record record; + + private RecordProjection(StructType structType, StructType projection) { + this(structType, projection, false); + } + + @SuppressWarnings("checkstyle:CyclomaticComplexity") + private RecordProjection(StructType structType, StructType projection, boolean allowMissing) { + this.type = projection; + this.positionMap = new int[projection.fields().size()]; + this.nestedProjections = new RecordProjection[projection.fields().size()]; + + // set up the projection positions and any nested projections that are needed + List dataFields = structType.fields(); + for (int pos = 0; pos < positionMap.length; pos += 1) { + Types.NestedField projectedField = projection.fields().get(pos); + + boolean found = false; + for (int i = 0; !found && i < dataFields.size(); i += 1) { + Types.NestedField dataField = dataFields.get(i); + if (projectedField.fieldId() == dataField.fieldId()) { + found = true; + positionMap[pos] = i; + switch (projectedField.type().typeId()) { + case STRUCT: + nestedProjections[pos] = + new RecordProjection( + dataField.type().asStructType(), projectedField.type().asStructType()); + break; + case MAP: + MapType projectedMap = projectedField.type().asMapType(); + MapType originalMap = dataField.type().asMapType(); + + boolean keyProjectable = + !projectedMap.keyType().isNestedType() + || projectedMap.keyType().equals(originalMap.keyType()); + boolean valueProjectable = + !projectedMap.valueType().isNestedType() + || projectedMap.valueType().equals(originalMap.valueType()); + Preconditions.checkArgument( + keyProjectable && valueProjectable, + "Cannot project a partial map key or value struct. Trying to project %s out of %s", + projectedField, + dataField); + + nestedProjections[pos] = null; + break; + case LIST: + ListType projectedList = projectedField.type().asListType(); + ListType originalList = dataField.type().asListType(); + + boolean elementProjectable = + !projectedList.elementType().isNestedType() + || projectedList.elementType().equals(originalList.elementType()); + Preconditions.checkArgument( + elementProjectable, + "Cannot project a partial list element struct. Trying to project %s out of %s", + projectedField, + dataField); + + nestedProjections[pos] = null; + break; + default: + nestedProjections[pos] = null; + } + } + } + + if (!found && projectedField.isOptional() && allowMissing) { + positionMap[pos] = -1; + nestedProjections[pos] = null; + } else if (!found) { + throw new IllegalArgumentException( + String.format("Cannot find field %s in %s", projectedField, structType)); + } + } + } + + public RecordProjection wrap(Record newRecord) { + this.record = newRecord; + return this; + } + + @Override + public int size() { + return type.fields().size(); + } + + @Override + public T get(int pos, Class javaClass) { + // struct can be null if wrap is not called first before the get call + // or if a null struct is wrapped. + if (record == null) { + return null; + } + + int recordPos = positionMap[pos]; + if (nestedProjections[pos] != null) { + Record nestedStruct = record.get(recordPos, Record.class); + if (nestedStruct == null) { + return null; + } + + return javaClass.cast(nestedProjections[pos].wrap(nestedStruct)); + } + + if (recordPos != -1) { + return record.get(recordPos, javaClass); + } else { + return null; + } + } + + @Override + public void set(int pos, T value) { + throw new UnsupportedOperationException(); + } + + @Override + public StructType struct() { + return type; + } + + @Override + public Object getField(String name) { + throw new UnsupportedOperationException(); + } + + @Override + public void setField(String name, Object value) { + throw new UnsupportedOperationException(); + } + + @Override + public Object get(int pos) { + return get(pos, Object.class); + } + + @Override + public Record copy() { + throw new UnsupportedOperationException(); + } + + @Override + public Record copy(Map overwriteValues) { + throw new UnsupportedOperationException(); + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/RecordWrapper.java b/core/src/main/java/kafka/automq/table/worker/RecordWrapper.java new file mode 100644 index 0000000000..9626d4fb92 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/RecordWrapper.java @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types.StructType; + +import java.util.Locale; +import java.util.Map; + +public class RecordWrapper implements Record { + + private final Record delegate; + private final Operation op; + + public RecordWrapper(Record delegate, String cdcField, boolean upsert) { + this.delegate = delegate; + String cdcValue = extract(delegate, cdcField); + if (cdcValue == null || cdcValue.isEmpty()) { + op = upsert ? Operation.UPDATE : Operation.INSERT; + } else { + switch (cdcValue.toUpperCase(Locale.ROOT).charAt(0)) { + case 'I': + op = Operation.INSERT; + break; + case 'U': + op = Operation.UPDATE; + break; + case 'D': + op = Operation.DELETE; + break; + default: + op = Operation.INSERT; + } + } + } + + public RecordWrapper(Record delegate, Operation op) { + this.delegate = delegate; + this.op = op; + } + + public Operation op() { + return op; + } + + @Override + public StructType struct() { + return delegate.struct(); + } + + @Override + public Object getField(String name) { + return delegate.getField(name); + } + + @Override + public void setField(String name, Object value) { + delegate.setField(name, value); + } + + @Override + public Object get(int pos) { + return delegate.get(pos); + } + + @Override + public Record copy() { + return new RecordWrapper(delegate.copy(), op); + } + + @Override + public Record copy(Map overwriteValues) { + return new RecordWrapper(delegate.copy(overwriteValues), op); + } + + @Override + public int size() { + return delegate.size(); + } + + @Override + public T get(int pos, Class javaClass) { + return delegate.get(pos, javaClass); + } + + @Override + public void set(int pos, T value) { + delegate.set(pos, value); + } + + public static String extract(Record record, String cdcField) { + if (cdcField == null) { + return null; + } + String[] fieldChain = cdcField.split("\\."); + return extract0(record, fieldChain, 0); + } + + public static String extract0(Object parent, String[] fieldChain, int idx) { + Object value = null; + if (parent instanceof Record) { + value = ((Record) parent).getField(fieldChain[idx]); + } else if (parent instanceof Map) { + value = ((Map) parent).get(fieldChain[idx]); + } + if (value == null) { + return null; + } + if (idx == fieldChain.length - 1) { + if (value instanceof String) { + return (String) value; + } else { + return null; + } + } else { + return extract0(value, fieldChain, idx + 1); + } + + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/RecordsIterator.java b/core/src/main/java/kafka/automq/table/worker/RecordsIterator.java new file mode 100644 index 0000000000..09ed3449e6 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/RecordsIterator.java @@ -0,0 +1,136 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import org.apache.kafka.clients.consumer.internals.CompletedFetch; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.record.ControlRecordType; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.RecordBatch; +import org.apache.kafka.common.utils.AbstractIterator; +import org.apache.kafka.common.utils.BufferSupplier; +import org.apache.kafka.common.utils.CloseableIterator; +import org.apache.kafka.storage.internals.log.FetchDataInfo; + +import java.util.Comparator; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.PriorityQueue; +import java.util.Set; + +/** + * ref. {@link CompletedFetch} + */ +public class RecordsIterator extends AbstractIterator { + private final Set abortedProducerIds = new HashSet<>(); + private final PriorityQueue abortedTransactions; + private final BufferSupplier bufferSupplier; + + private final Iterator batches; + private CloseableIterator records; + private long nextOffset; + + public RecordsIterator(long startOffset, FetchDataInfo rst, BufferSupplier bufferSupplier) { + this.nextOffset = startOffset; + this.batches = rst.records.batchIterator(); + this.abortedTransactions = abortedTransactions(rst.abortedTransactions.orElse(null)); + this.bufferSupplier = bufferSupplier; + } + + public long nextOffset() { + return nextOffset; + } + + @Override + protected Record makeNext() { + if (records != null && records.hasNext()) { + Record record = records.next(); + nextOffset = record.offset() + 1; + return record; + } + + while (batches.hasNext()) { + if (records != null) { + records.close(); + records = null; + } + RecordBatch currentBatch = batches.next(); + if (currentBatch.hasProducerId()) { + consumeAbortedTransactionsUpTo(currentBatch.lastOffset()); + long producerId = currentBatch.producerId(); + if (containsAbortMarker(currentBatch)) { + abortedProducerIds.remove(producerId); + } else if (isBatchAborted(currentBatch)) { + nextOffset = currentBatch.nextOffset(); + continue; + } + } + if (currentBatch.isControlBatch()) { + nextOffset = currentBatch.nextOffset(); + continue; + } + records = currentBatch.streamingIterator(bufferSupplier); + return makeNext(); + } + if (records != null) { + records.close(); + } + return allDone(); + } + + private void consumeAbortedTransactionsUpTo(long offset) { + if (abortedTransactions == null) + return; + + while (!abortedTransactions.isEmpty() && abortedTransactions.peek().firstOffset() <= offset) { + FetchResponseData.AbortedTransaction abortedTransaction = abortedTransactions.poll(); + abortedProducerIds.add(abortedTransaction.producerId()); + } + } + + private boolean isBatchAborted(RecordBatch batch) { + return batch.isTransactional() && abortedProducerIds.contains(batch.producerId()); + } + + private PriorityQueue abortedTransactions( + List abortedTransactionList) { + if (abortedTransactionList == null || abortedTransactionList.isEmpty()) + return null; + + PriorityQueue abortedTransactions = new PriorityQueue<>( + abortedTransactionList.size(), Comparator.comparingLong(FetchResponseData.AbortedTransaction::firstOffset) + ); + abortedTransactions.addAll(abortedTransactionList); + return abortedTransactions; + } + + private boolean containsAbortMarker(RecordBatch batch) { + if (!batch.isControlBatch()) + return false; + + Iterator batchIterator = batch.iterator(); + if (!batchIterator.hasNext()) + return false; + + Record firstRecord = batchIterator.next(); + return ControlRecordType.ABORT == ControlRecordType.parse(firstRecord.key()); + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/ReferenceHolder.java b/core/src/main/java/kafka/automq/table/worker/ReferenceHolder.java new file mode 100644 index 0000000000..b113eae380 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/ReferenceHolder.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import io.netty.util.AbstractReferenceCounted; +import io.netty.util.ReferenceCounted; + +public class ReferenceHolder extends AbstractReferenceCounted { + private final T value; + private final Runnable deallocate; + + public ReferenceHolder(T value, Runnable deallocate) { + this.value = value; + this.deallocate = deallocate; + } + + public T value() { + return value; + } + + @Override + protected void deallocate() { + deallocate.run(); + } + + @Override + public ReferenceCounted touch(Object o) { + return this; + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/TableWorkers.java b/core/src/main/java/kafka/automq/table/worker/TableWorkers.java new file mode 100644 index 0000000000..b26d066139 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/TableWorkers.java @@ -0,0 +1,163 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.Channel; +import kafka.automq.table.events.CommitRequest; +import kafka.automq.table.events.Envelope; +import kafka.automq.table.process.RecordProcessorFactory; +import kafka.automq.table.utils.TableIdentifierUtil; +import kafka.cluster.Partition; +import kafka.server.KafkaConfig; + +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.ThreadUtils; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.threads.EventLoop; + +import org.apache.iceberg.catalog.Catalog; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.Semaphore; + +public class TableWorkers { + private static final Logger LOGGER = LoggerFactory.getLogger(TableWorkers.class); + public static final ScheduledExecutorService SCHEDULER = Threads.newSingleThreadScheduledExecutor("table-workers", true, LOGGER); + private final Catalog catalog; + private final Channel channel; + private final Channel.SubChannel subChannel; + private final EventLoopWorker[] workers; + private final EventLoops eventLoops; + private final ExecutorService executor = Threads.newFixedThreadPool(1, ThreadUtils.createThreadFactory("table-worker-poll", true), LOGGER); + private final ExecutorService flushExecutor = Threads.newFixedThreadPool(Systems.CPU_CORES, ThreadUtils.createThreadFactory("table-workers-flush", true), LOGGER); + private final KafkaConfig config; + private final RecordProcessorFactory recordProcessorFactory; + private final Semaphore commitLimiter = new Semaphore(Systems.CPU_CORES); + private volatile boolean closed = false; + + + public TableWorkers(Catalog catalog, Channel channel, KafkaConfig config) { + this.catalog = catalog; + this.channel = channel; + this.subChannel = channel.subscribeControl(); + workers = new EventLoopWorker[Math.max(Systems.CPU_CORES / 2, 1)]; + EventLoop[] eventLoops = new EventLoop[workers.length]; + for (int i = 0; i < workers.length; i++) { + workers[i] = new EventLoopWorker(i); + eventLoops[i] = workers[i].eventLoop; + } + this.eventLoops = new EventLoops(eventLoops); + executor.submit(new ControlListener()); + this.config = config; + this.recordProcessorFactory = new RecordProcessorFactory(config.tableTopicSchemaRegistryUrl()); + } + + public void add(Partition partition) { + workers[Math.abs(partition.topic().hashCode() % workers.length)].add(partition); + } + + public void remove(Partition partition) { + workers[Math.abs(partition.topic().hashCode() % workers.length)].remove(partition); + } + + public synchronized void close() { + closed = true; + this.subChannel.close(); + } + + class EventLoopWorker { + private final EventLoop eventLoop; + private final Map topic2worker = new ConcurrentHashMap<>(); + + public EventLoopWorker(int index) { + eventLoop = new EventLoop("table-worker-" + index); + } + + public void add(Partition partition) { + eventLoop.execute(() -> { + topic2worker.compute(partition.topic(), (topic, worker) -> { + if (worker == null) { + WorkerConfig config = new WorkerConfig(partition); + IcebergWriterFactory writerFactory = new IcebergWriterFactory(catalog, + TableIdentifierUtil.of(config.namespace(), partition.topic()), recordProcessorFactory, config, partition.topic()); + worker = new TopicPartitionsWorker(partition.topic(), config, + writerFactory, channel, eventLoop, eventLoops, flushExecutor, commitLimiter); + } + worker.add(partition); + return worker; + }); + }); + } + + public void remove(Partition partition) { + eventLoop.execute(() -> { + TopicPartitionsWorker topicPartitionsWorker = topic2worker.get(partition.topic()); + if (topicPartitionsWorker == null) { + return; + } + topicPartitionsWorker.remove(partition); + if (topicPartitionsWorker.isEmpty()) { + topic2worker.remove(partition.topic()); + } + }); + } + + public void commit(CommitRequest commitRequest) { + if (topic2worker.containsKey(commitRequest.topic())) { + eventLoop.execute(() -> { + TopicPartitionsWorker topicPartitionsWorker = topic2worker.get(commitRequest.topic()); + if (topicPartitionsWorker != null) { + topicPartitionsWorker.commit(commitRequest); + } + }); + } + } + } + + class ControlListener implements Runnable { + + @Override + public void run() { + for (; ; ) { + Envelope envelope; + synchronized (TableWorkers.this) { + if (closed) { + return; + } + envelope = subChannel.poll(); + } + if (envelope == null) { + Threads.sleep(10); + continue; + } + CommitRequest commitRequest = envelope.event().payload(); + for (EventLoopWorker worker : workers) { + worker.commit(commitRequest); + } + } + } + } + +} diff --git a/core/src/main/java/kafka/automq/table/worker/TopicPartitionsWorker.java b/core/src/main/java/kafka/automq/table/worker/TopicPartitionsWorker.java new file mode 100644 index 0000000000..6f61b2b8aa --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/TopicPartitionsWorker.java @@ -0,0 +1,783 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.Channel; +import kafka.automq.table.events.CommitRequest; +import kafka.automq.table.events.CommitResponse; +import kafka.automq.table.events.Errors; +import kafka.automq.table.events.Event; +import kafka.automq.table.events.EventType; +import kafka.automq.table.events.PartitionMetric; +import kafka.automq.table.events.TopicMetric; +import kafka.automq.table.events.WorkerOffset; +import kafka.cluster.Partition; +import kafka.cluster.PartitionAppendListener; +import kafka.log.UnifiedLog; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.storage.internals.log.LogSegment; + +import com.automq.stream.s3.metrics.TimerUtil; +import com.automq.stream.utils.threads.EventLoop; + +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.io.WriteResult; +import org.apache.iceberg.types.Types; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; + +/** + * The worker manage the topic partitions iceberg sync. + *

+ * There are two paths to sync the records to iceberg: + * 1. Advance sync: The worker will sync the 'hot data' to the writer when the dirty bytes reach the threshold. + * 2. Commit sync: The worker receive {@link CommitRequest}, sync the records [start offset, log end) to iceberg + * and response the result. + *

+ * Thread-Safe: the event loop sequences all operations on the worker. + */ +class TopicPartitionsWorker { + private static final Logger LOGGER = LoggerFactory.getLogger(TopicPartitionsWorker.class); + private static final int WRITER_LIMIT = 1000; + private static final AtomicLong SYNC_ID = new AtomicLong(0); + + private final Semaphore commitLimiter; + + private final String topic; + private final String name; + private final WorkerConfig config; + + Status status = Status.IDLE; + + private final List commitRequests = new LinkedList<>(); + + List writers = new ArrayList<>(); + + /** + * The flag to indicate the worker need reset. The worker status(writers) will be reset in the next round run. + */ + private boolean requireReset; + + private Map partitions = new HashMap<>(); + + private double avgRecordSize = Double.MIN_VALUE; + private double decompressedRatio = 1.0; + + final Commit commit = new Commit(); + final AdvanceSync advanceSync = new AdvanceSync(); + + private final EventLoop eventLoop; + private final EventLoops eventLoops; + private final ExecutorService flushExecutors; + private final Channel channel; + private final WriterFactory writerFactory; + + public TopicPartitionsWorker(String topic, WorkerConfig config, WriterFactory writerFactory, Channel channel, + EventLoop eventLoop, EventLoops eventLoops, ExecutorService flushExecutors, Semaphore commitLimiter) { + this.topic = topic; + this.name = topic; + this.config = config; + this.eventLoop = eventLoop; + this.eventLoops = eventLoops; + this.channel = channel; + this.writerFactory = writerFactory; + this.flushExecutors = flushExecutors; + this.commitLimiter = commitLimiter; + } + + public void add(Partition partition) { + Map newPartitions = new HashMap<>(partitions); + newPartitions.put(partition.partitionId(), partition); + partition.addAppendListener(advanceSync); + this.partitions = newPartitions; + requireReset = true; + LOGGER.info("[ADD_PARTITION],{},{}", name, partition.partitionId()); + } + + public void remove(Partition partition) { + Map newPartitions = new HashMap<>(partitions); + newPartitions.remove(partition.partitionId()); + partition.removeAppendListener(advanceSync); + this.partitions = newPartitions; + requireReset = true; + LOGGER.info("[REMOVE_PARTITION],{},{}", name, partition.partitionId()); + } + + public CompletableFuture commit(CommitRequest commitRequest) { + RequestWrapper request = new RequestWrapper(commitRequest); + commitRequests.add(request); + run(); + return request.cf; + } + + public Boolean isEmpty() { + return partitions.isEmpty(); + } + + void run() { + if (status != Status.IDLE) { + return; + } + if (partitions.isEmpty()) { + LOGGER.info("[WORKER_EMPTY],{}", topic); + return; + } + if (requireReset) { + writers.clear(); + requireReset = false; + } + try { + if (commitRequests.isEmpty() && !commit.starving) { + advanceSync.sync(); + } else { + commit.commit(); + } + } catch (Throwable e) { + transitionTo(Status.IDLE); + requireReset = true; + LOGGER.error("[WORKER_RUN_FAIL],{}", topic, e); + } + } + + private boolean transitionTo(Status status) { + if (this.status == status) { + return true; + } + if (status == Status.IDLE) { + this.status = status; + commitLimiter.release(); + return true; + } + if (commitLimiter.tryAcquire()) { + this.status = status; + return true; + } else { + return false; + } + } + + private ReferenceHolder tryAcquireMoreCommitQuota(int required) { + for (; ; ) { + int available = commitLimiter.availablePermits(); + int real = Math.min(available, required); + if (commitLimiter.tryAcquire(real)) { + return new ReferenceHolder<>(real, () -> commitLimiter.release(real)); + } + } + } + + /** + * Remove commited writers and clear useless writer. + */ + static void cleanupCommitedWriters(String topic, List offsets, + List writers) throws IOException { + int matchIndex = -1; + for (int i = 0; i < writers.size(); i++) { + if (isOffsetMatch(writers.get(i), offsets)) { + matchIndex = i; + } + } + // clear useless writer + if (matchIndex == -1) { + if (!writers.isEmpty()) { + Writer lastWriter = writers.get(writers.size() - 1); + if (!lastWriter.isCompleted()) { + writers.get(writers.size() - 1).abort(); + LOGGER.info("[WRITER_NOT_MATCH],[RESET],{}", topic); + } + } + writers.clear(); + } else { + List newWriters = new ArrayList<>(writers.subList(matchIndex, writers.size())); + writers.clear(); + writers.addAll(newWriters); + } + cleanupEmptyWriter(writers); + } + + static void cleanupEmptyWriter(List writers) { + int size = writers.size(); + Iterator it = writers.iterator(); + while (size > 0 && it.hasNext()) { + Writer writer = it.next(); + if (writer.isCompleted() && writer.results().isEmpty()) { + it.remove(); + size--; + } + } + } + + static boolean isOffsetMatch(Writer writer, List request) { + for (WorkerOffset workerOffset : request) { + Writer.OffsetRange writerOffset = writer.getOffset(workerOffset.partition()); + if (writerOffset == null || writerOffset.start() != workerOffset.offset()) { + return false; + } + } + return true; + } + + protected double getAvgRecordSize() { + for (Map.Entry entry : partitions.entrySet()) { + Partition partition = entry.getValue(); + UnifiedLog log = partition.log().get(); + LogSegment segment = log.activeSegment(); + long recordCount = log.highWatermark() - segment.baseOffset(); + if (recordCount <= 0) { + continue; + } + avgRecordSize = ((double) segment.size()) / recordCount; + } + return avgRecordSize; + } + + protected double getDecompressedRatio() { + return decompressedRatio; + } + + protected void updateDecompressedRatio(double decompressedRatio) { + this.decompressedRatio = decompressedRatio; + } + + SyncTask newSyncTask(String type, Map startOffsets, long priority) { + return new SyncTask(String.format("[%s-%s-%s]", type, topic, SYNC_ID.getAndIncrement()), partitions, startOffsets, getAvgRecordSize(), priority); + } + + class SyncTask { + private final String logContext; + private final Map partitions; + private final Map startOffsets; + private boolean hasMoreData; + private ReferenceHolder moreQuota; + private final List writers = new ArrayList<>(); + private final double avgRecordSize; + private final double decompressedRatio; + private final long priority; + + final List microSyncBatchTasks = new ArrayList<>(); + + public SyncTask(String logContext, Map partitions, Map startOffsets, + double avgRecordSize, long taskPriority) { + this.logContext = logContext; + this.partitions = partitions; + this.startOffsets = new HashMap<>(startOffsets); + this.avgRecordSize = avgRecordSize; + this.priority = taskPriority; + this.decompressedRatio = getDecompressedRatio(); + } + + public boolean hasMoreData() { + return hasMoreData; + } + + /** + * Split sync task to parallel micro sync batch tasks. + */ + // visible to test + void plan() { + Map partition2bound = new HashMap<>(); + for (Map.Entry entry : startOffsets.entrySet()) { + int partitionId = entry.getKey(); + Partition partition = partitions.get(partitionId); + OffsetBound offsetBound = new OffsetBound(partition, startOffsets.get(partitionId), partition.log().get().highWatermark()); + partition2bound.put(partitionId, offsetBound); + } + double totalDirtyBytes = partition2bound.values().stream().mapToLong(bound -> bound.end - bound.start).sum() * avgRecordSize * decompressedRatio; + int parallel = Math.max((int) (Math.ceil(totalDirtyBytes / config.microSyncBatchSize())), 1); + parallel = Math.min(Math.min(parallel, eventLoops.size()), partition2bound.size()); + moreQuota = tryAcquireMoreCommitQuota(parallel - 1); + parallel = moreQuota.value() + 1; + + List partitionSortedBySize = partition2bound.entrySet().stream().sorted((e1, e2) -> { + long size1 = e1.getValue().end - e1.getValue().start; + long size2 = e2.getValue().end - e2.getValue().start; + return Long.compare(size2, size1); + }).map(Map.Entry::getKey).collect(Collectors.toList()); + + for (int i = 0; i < parallel; i++) { + EventLoops.EventLoopRef workerEventLoop = eventLoops.leastLoadEventLoop(); + Writer taskWriter = writerFactory.newWriter(); + writers.add(taskWriter); + PartitionWriteTaskContext ctx = new PartitionWriteTaskContext(taskWriter, workerEventLoop, flushExecutors, config, priority); + microSyncBatchTasks.add(new MicroSyncBatchTask(logContext, ctx)); + } + for (int i = 0; i < partitionSortedBySize.size(); i++) { + int partitionId = partitionSortedBySize.get(i); + microSyncBatchTasks.get(i % parallel).addPartition(partition2bound.get(partitionId)); + } + for (MicroSyncBatchTask task : microSyncBatchTasks) { + task.startOffsets(startOffsets); + task.endOffsets(config.microSyncBatchSize(), avgRecordSize * decompressedRatio); + task.offsetBounds().forEach(bound -> startOffsets.put(bound.partition.partitionId(), bound.end)); + if (task.hasMoreData()) { + hasMoreData = true; + } + } + } + + public CompletableFuture sync() { + TimerUtil timer = new TimerUtil(); + plan(); + return CompletableFuture.allOf(microSyncBatchTasks + .stream() + .map(t -> t.run().thenApply(rst -> { + if (rst.ctx.requireReset) { + TopicPartitionsWorker.this.requireReset = true; + } + updateDecompressedRatio(); + return rst; + })).toArray(CompletableFuture[]::new) + ) + .thenApply(nil -> this) + .whenComplete((rst, ex) -> { + moreQuota.release(); + if (ex != null) { + LOGGER.error("[SYNC_TASK_FAIL],{}", logContext, ex); + TopicPartitionsWorker.this.requireReset = true; + } else { + TopicPartitionsWorker.this.writers.addAll(writers); + List offsetBounds = microSyncBatchTasks.stream().flatMap(t -> t.offsetBounds().stream()).collect(Collectors.toList()); + double size = avgRecordSize * offsetBounds.stream().mapToLong(bound -> bound.end - bound.start).sum(); + LOGGER.info("[SYNC_TASK],{},size={},decompressedRatio={},parallel={},partitions={},cost={}ms", logContext, size, decompressedRatio, microSyncBatchTasks.size(), offsetBounds, timer.elapsedAs(TimeUnit.MILLISECONDS)); + } + }); + } + + private void updateDecompressedRatio() { + MicroSyncBatchTask task = microSyncBatchTasks.get(0); + double recordBatchSize = task.offsetBounds.stream().mapToLong(bound -> bound.end - bound.start).sum() * avgRecordSize; + if (recordBatchSize == 0) { + return; + } + TopicPartitionsWorker.this.updateDecompressedRatio(Math.max(task.ctx.writeSize / recordBatchSize, 1.0)); + } + } + + /** + * The micro sync batch task syncs the records in partition - [start, end) to iceberg. + * How to use: + * 1. #addPartition: add the partition to the task. + * 2. #startOffsets: set start offsets for partition syncing. + * 3. #endOffset: adjust the end offsets for partition syncing according to the batch size and record size. + * 4. #run: run the task. + */ + static class MicroSyncBatchTask { + private final List offsetBounds = new ArrayList<>(); + private boolean hasMoreData; + private final String logContext; + private final PartitionWriteTaskContext ctx; + + public MicroSyncBatchTask(String logContext, PartitionWriteTaskContext ctx) { + this.logContext = logContext; + this.ctx = ctx; + } + + public List offsetBounds() { + return offsetBounds; + } + + public boolean hasMoreData() { + return hasMoreData; + } + + public void addPartition(OffsetBound offsetBound) { + offsetBounds.add(offsetBound); + } + + public void startOffsets(Map startOffsets) { + offsetBounds.forEach(bound -> bound.start = startOffsets.get(bound.partition.partitionId())); + startOffsets.forEach(ctx.writer::setOffset); + } + + public void endOffsets(int microSyncBatchSize, double avgRecordSize) { + int avgRecordCountPerPartition = Math.max((int) Math.ceil(microSyncBatchSize / avgRecordSize / offsetBounds.size()), 1); + long[] fixedEndOffsets = new long[offsetBounds.size()]; + hasMoreData = offsetBounds.stream().mapToLong(bound -> bound.end - bound.start).sum() * avgRecordSize > microSyncBatchSize; + if (!hasMoreData) { + return; + } + int remaining = 0; + for (int i = 0; i < offsetBounds.size(); i++) { + OffsetBound offsetBound = offsetBounds.get(i); + int count = (int) Math.min(offsetBound.end - offsetBound.start, avgRecordCountPerPartition); + fixedEndOffsets[i] = offsetBound.start + count; + remaining += avgRecordCountPerPartition - count; + } + for (int i = 0; i < offsetBounds.size() && remaining > 0; i++) { + OffsetBound offsetBound = offsetBounds.get(i); + int add = (int) Math.min(offsetBound.end - fixedEndOffsets[i], remaining); + fixedEndOffsets[i] += add; + remaining -= add; + } + for (int i = 0; i < offsetBounds.size(); i++) { + offsetBounds.get(i).end = fixedEndOffsets[i]; + } + } + + public CompletableFuture run() { + List> partitionCfList = new ArrayList<>(); + offsetBounds.forEach(bound -> partitionCfList.add(runPartitionWriteTask(bound.partition, bound.start, bound.end))); + return CompletableFuture.allOf(partitionCfList.toArray(new CompletableFuture[0])) + .thenComposeAsync(nil -> { + if (ctx.requireReset) { + return CompletableFuture.completedFuture(null); + } else { + return ctx.writer.flush(FlushMode.COMPLETE, ctx.flushExecutors, ctx.eventLoop) + .exceptionally(ex -> { + ctx.requireReset = true; + LOGGER.error("[FLUSH_FAIL],{}", logContext, ex); + return null; + }); + } + }, ctx.eventLoop).thenApply(nil -> this).whenComplete((nil, ex) -> { + ctx.eventLoop.release(); + }); + } + + protected CompletableFuture runPartitionWriteTask(Partition partition, long start, long end) { + return new PartitionWriteTask(partition, start, end, ctx).run(); + } + } + + class AdvanceSync implements PartitionAppendListener { + /** + * The dirty bytes which partition appends but not sync to the writer. It's used to trigger the advanced sync. + */ + private final AtomicLong dirtyBytesCounter = new AtomicLong(0); + private final AtomicBoolean fastNextSync = new AtomicBoolean(false); + + /** + * Try to sync the 'hot data' to the writer. + *

+ * Note: it should run in event loop. + */ + public void sync() { + Optional> startOffsetsOpt = getStartOffsets(); + if (startOffsetsOpt.isEmpty()) { + return; + } + fastNextSync.set(false); + SyncTask syncTask = newSyncTask("ADVANCED", startOffsetsOpt.get(), System.currentTimeMillis()); + syncTask.sync().whenCompleteAsync((task, e) -> { + transitionTo(Status.IDLE); + fastNextSync.set(task.hasMoreData()); + run(); + }, eventLoop); + + } + + @Override + public void onAppend(TopicPartition partition, MemoryRecords records) { + if (dirtyBytesCounter.addAndGet(records.sizeInBytes()) * decompressedRatio > config.incrementSyncThreshold()) { + if (fastNextSync.compareAndSet(false, true)) { + eventLoop.submit(TopicPartitionsWorker.this::run); + } + } + } + + private Optional> getStartOffsets() { + if (!transitionTo(Status.ADVANCE_SYNC)) { + TableWorkers.SCHEDULER.schedule(() -> eventLoop.submit(TopicPartitionsWorker.this::run), 100, TimeUnit.MILLISECONDS); + return Optional.empty(); + } + if (!fastNextSync.get() + // The commit speed can't keep up with the advance sync rate. + // We need to limit the number of writers to prevent excessive heap usage. + || writers.size() >= WRITER_LIMIT + ) { + transitionTo(Status.IDLE); + return Optional.empty(); + } + if (writers.isEmpty()) { + // Await for the first commit request to know the sync start offset + transitionTo(Status.IDLE); + return Optional.empty(); + } + Writer lastWriter = writers.get(writers.size() - 1); + Map startOffsets = new HashMap<>(); + lastWriter.getOffsets().forEach((partition, offset) -> startOffsets.put(partition, offset.end())); + this.dirtyBytesCounter.set(0); + return Optional.of(startOffsets); + } + + } + + class Commit { + boolean starving = false; + private List lastCommitRequestOffsets = Collections.emptyList(); + + /** + * Handle the pending commit requests. + *

+ * Note: it should run in event loop. + */ + void commit() throws Exception { + if (!transitionTo(Status.COMMIT)) { + TableWorkers.SCHEDULER.schedule(() -> eventLoop.submit(TopicPartitionsWorker.this::run), 1, TimeUnit.SECONDS); + return; + } + RequestWrapper requestWrapper = preCommit(); + if (requestWrapper == null) { + transitionTo(Status.IDLE); + return; + } + if (writerFactory.partitionSpec() != null && requestWrapper.request.specId() > writerFactory.partitionSpec().specId()) { + writerFactory.reset(); + writers.clear(); + LOGGER.info("[WORKER_APPLY_NEW_PARTITION],{},spec={}", this, config.partitionByConfig()); + } + config.refresh(); + + cleanupCommitedWriters(requestWrapper.request.offsets()); + if (requestWrapper.request.offsets().equals(lastCommitRequestOffsets) && !writers.isEmpty()) { + // fast commit previous timeout request + commitResponse(requestWrapper, true); + starving = true; + return; + } + starving = false; + advanceSync.dirtyBytesCounter.set(0); + Writer lastWriter = writers.isEmpty() ? null : writers.get(writers.size() - 1); + Map startOffsets = requestWrapper.request.offsets().stream().collect(Collectors.toMap( + WorkerOffset::partition, + o -> Optional.ofNullable(lastWriter) + .map(w -> w.getOffsets().get(o.partition())) + .map(offsetRange -> offsetRange.end) + .orElse(o.offset()) + )); + SyncTask syncTask = newSyncTask("COMMIT", startOffsets, requestWrapper.timestamp); + syncTask.sync().whenCompleteAsync((task, e) -> { + if (e == null) { + lastCommitRequestOffsets = requestWrapper.request.offsets(); + } + commitResponse(requestWrapper, task.hasMoreData()); + if (task.hasMoreData()) { + advanceSync.fastNextSync.set(true); + } + }, eventLoop); + } + + /** + * 1. Remove outdated commit requests. + * 2. Filter the partitions need resolved by current worker and EPOCH_MISMATCH response with whose + * partition epoch not equals to request epoch. + * 3. Generate a new {@link RequestWrapper} with remaining valid partitions. + * + * @return {@link RequestWrapper} the request needs to be handled, or null if there is no request needs to be handled. + */ + private RequestWrapper preCommit() throws Exception { + if (commitRequests.isEmpty()) { + return null; + } + if (commitRequests.size() > 1) { + LOGGER.warn("[DROP_COMMIT_REQUEST],{}", commitRequests.subList(0, commitRequests.size() - 1)); + } + RequestWrapper requestWrapper = commitRequests.get(commitRequests.size() - 1); + commitRequests.clear(); + + CommitRequest request = requestWrapper.request; + List epochMismatchPartitions = new ArrayList<>(); + List requestCommitPartitions = new ArrayList<>(); + // filter the partitions need resolved by current worker. + for (WorkerOffset offset : request.offsets()) { + Partition partition = partitions.get(offset.partition()); + if (partition == null) { + continue; + } + if (partition.getLeaderEpoch() != offset.epoch()) { + epochMismatchPartitions.add(partition); + } else { + requestCommitPartitions.add(offset); + } + } + if (!epochMismatchPartitions.isEmpty()) { + CommitResponse commitResponse = mismatchEpochResponse(request, epochMismatchPartitions); + channel.asyncSend(topic, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, commitResponse)); + LOGGER.info("[COMMIT_RESPONSE],[EPOCH_MISMATCH],{}", commitResponse); + } + if (requestCommitPartitions.isEmpty()) { + return null; + } + RequestWrapper wrapper = new RequestWrapper(new CommitRequest(request.commitId(), request.topic(), request.specId(), requestCommitPartitions), requestWrapper.cf, requestWrapper.timestamp); + if (System.currentTimeMillis() - wrapper.timestamp >= TimeUnit.SECONDS.toMillis(30)) { + LOGGER.warn("[DROP_EXPIRED_COMMIT_REQUEST],{}", wrapper.request.commitId()); + starving = true; + return null; + } + return wrapper; + } + + private CompletableFuture commitResponse(RequestWrapper requestWrapper, boolean fastCommit) { + CompletableFuture cf = CompletableFuture.completedFuture(null); + cf = cf.thenComposeAsync(nil -> { + if (requireReset) { + CommitResponse commitResponse = new CommitResponse(Types.StructType.of(), Errors.MORE_DATA, requestWrapper.request.commitId(), + topic, requestWrapper.request.offsets(), Collections.emptyList(), Collections.emptyList(), TopicMetric.NOOP, Collections.emptyList()); + return channel.asyncSend(topic, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, commitResponse)) + .thenAccept(rst -> LOGGER.info("[COMMIT_RESPONSE],{}", commitResponse)); + } + Map partition2epoch = requestWrapper.request.offsets().stream().collect(Collectors.toMap(WorkerOffset::partition, WorkerOffset::epoch)); + Map partitionNextOffsetMap = new HashMap<>(); + + List dataFiles = new ArrayList<>(); + List deleteFiles = new ArrayList<>(); + Map partitionMetricMap = new HashMap<>(); + long fieldCount = 0; + for (Writer writer : writers) { + List writeResults = writer.results(); + writeResults.forEach(writeResult -> { + dataFiles.addAll(List.of(writeResult.dataFiles())); + deleteFiles.addAll(List.of(writeResult.deleteFiles())); + }); + writer.getOffsets().forEach((partition, offsetRange) -> partitionNextOffsetMap.put(partition, offsetRange.end())); + writer.partitionMetrics().forEach((partition, metric) -> partitionMetricMap.compute(partition, (p, oldMetric) -> { + if (oldMetric == null) { + return metric; + } else { + return metric.watermark() > oldMetric.watermark() ? metric : oldMetric; + } + })); + fieldCount += writer.topicMetric().fieldCount(); + } + List nextOffsets = partitionNextOffsetMap.entrySet().stream() + .map(e -> new WorkerOffset(e.getKey(), partition2epoch.get(e.getKey()), e.getValue())) + .collect(Collectors.toList()); + List partitionMetrics = partitionMetricMap.entrySet().stream() + .map(e -> new PartitionMetric(e.getKey(), e.getValue().watermark())) + .collect(Collectors.toList()); + TopicMetric topicMetric = new TopicMetric(fieldCount); + + PartitionSpec spec = writerFactory.partitionSpec(); + Types.StructType partitionType = spec != null ? spec.partitionType() : Types.StructType.of(); + CommitResponse commitResponse = new CommitResponse(partitionType, + fastCommit ? Errors.MORE_DATA : Errors.NONE, + requestWrapper.request.commitId(), topic, nextOffsets, dataFiles, deleteFiles, topicMetric, partitionMetrics); + return channel.asyncSend(topic, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, commitResponse)) + .thenAccept(rst -> LOGGER.info("[COMMIT_RESPONSE],{}", commitResponse)); + }, eventLoop).whenCompleteAsync((nil, ex) -> { + if (ex != null) { + requireReset = true; + LOGGER.error("[COMMIT_RESPONSE_FAIL],{}", topic, ex); + requestWrapper.cf.completeExceptionally(ex); + } else { + requestWrapper.cf.complete(null); + } + transitionTo(Status.IDLE); + run(); + }, eventLoop); + return cf; + } + + private void cleanupCommitedWriters(List offsets) throws IOException { + TopicPartitionsWorker.cleanupCommitedWriters(topic, offsets, writers); + } + } + + static CommitResponse mismatchEpochResponse(CommitRequest request, List partitions) { + return new CommitResponse( + Types.StructType.of(), + Errors.EPOCH_MISMATCH, + request.commitId(), + request.topic(), + partitions.stream().map(p -> new WorkerOffset(p.partitionId(), p.getLeaderEpoch(), -1)).collect(Collectors.toList()), + Collections.emptyList(), + Collections.emptyList(), + TopicMetric.NOOP, + Collections.emptyList() + ); + } + + static class RequestWrapper { + CommitRequest request; + // only used for test + CompletableFuture cf = new CompletableFuture<>(); + long timestamp; + + public RequestWrapper(CommitRequest request) { + this.request = request; + this.timestamp = System.currentTimeMillis(); + } + + public RequestWrapper(CommitRequest request, CompletableFuture cf, long timestamp) { + this.request = request; + this.cf = cf; + this.timestamp = timestamp; + } + } + + enum Status { + ADVANCE_SYNC, + COMMIT, + IDLE, + } + + static class OffsetBound { + final Partition partition; + long start; + long end; + + public OffsetBound(Partition partition, long start, long end) { + this.partition = partition; + this.start = start; + this.end = end; + } + + @Override + public String toString() { + return String.format("%s-%s-%s", partition.partitionId(), start, end); + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + OffsetBound that = (OffsetBound) o; + return start == that.start && end == that.end && Objects.equals(partition, that.partition); + } + + @Override + public int hashCode() { + return Objects.hash(partition, start, end); + } + } + +} diff --git a/core/src/main/java/kafka/automq/table/worker/UnpartitionedDeltaWriter.java b/core/src/main/java/kafka/automq/table/worker/UnpartitionedDeltaWriter.java new file mode 100644 index 0000000000..3054b893d3 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/UnpartitionedDeltaWriter.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package kafka.automq.table.worker; + +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.OutputFileFactory; + +import java.io.IOException; +import java.util.Set; + +public class UnpartitionedDeltaWriter extends BaseDeltaTaskWriter { + private final RowDataDeltaWriter writer; + + UnpartitionedDeltaWriter( + PartitionSpec spec, + FileFormat format, + FileAppenderFactory appenderFactory, + OutputFileFactory fileFactory, + FileIO io, + long targetFileSize, + Schema schema, + Set identifierFieldIds + ) { + super( + spec, + format, + appenderFactory, + fileFactory, + io, + targetFileSize, + schema, + identifierFieldIds); + this.writer = new RowDataDeltaWriter(null); + } + + @Override + RowDataDeltaWriter partition(Record row) { + return writer; + } + + @Override + public void close() throws IOException { + writer.close(); + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/WorkerConfig.java b/core/src/main/java/kafka/automq/table/worker/WorkerConfig.java new file mode 100644 index 0000000000..1c297a66e3 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/WorkerConfig.java @@ -0,0 +1,128 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.cluster.Partition; +import kafka.log.UnifiedLog; + +import org.apache.kafka.common.config.TopicConfig; +import org.apache.kafka.server.record.ErrorsTolerance; +import org.apache.kafka.server.record.TableTopicConvertType; +import org.apache.kafka.server.record.TableTopicSchemaType; +import org.apache.kafka.server.record.TableTopicTransformType; +import org.apache.kafka.storage.internals.log.LogConfig; + +import java.util.List; + +import static org.apache.kafka.server.common.automq.TableTopicConfigValidator.stringToList; + +public class WorkerConfig { + static final String COMMA_NO_PARENS_REGEX = ",(?![^()]*+\\))"; + + private final UnifiedLog log; + private LogConfig config; + + public WorkerConfig(Partition partition) { + this.log = partition.log().get(); + this.config = log.config(); + } + + // for testing + public WorkerConfig() { + this.log = null; + } + + public String namespace() { + return config.tableTopicNamespace; + } + + public TableTopicSchemaType schemaType() { + return config.tableTopicSchemaType; + } + + public TableTopicConvertType valueConvertType() { + return config.valueConvertType; + } + public TableTopicConvertType keyConvertType() { + return config.keyConvertType; + } + + public String valueSubject() { + return config.getString(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_SUBJECT_CONFIG); + } + + public String valueMessageFullName() { + return config.getString(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_CONFIG); + } + + public String keySubject() { + return config.getString(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_SUBJECT_CONFIG); + } + + public String keyMessageFullName() { + return config.getString(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_CONFIG); + } + + public TableTopicTransformType transformType() { + return config.transformType; + } + + + public long incrementSyncThreshold() { + return 32 * 1024 * 1024; + } + + + public int microSyncBatchSize() { + return 32 * 1024 * 1024; + } + + public List idColumns() { + String str = config.tableTopicIdColumns; + return stringToList(str, COMMA_NO_PARENS_REGEX); + } + + public String partitionByConfig() { + return config.tableTopicPartitionBy; + } + + public List partitionBy() { + String str = config.tableTopicPartitionBy; + return stringToList(str, COMMA_NO_PARENS_REGEX); + } + + public boolean upsertEnable() { + return config.tableTopicUpsertEnable; + } + + public String cdcField() { + return config.tableTopicCdcField; + } + + public ErrorsTolerance errorsTolerance() { + return config.errorsTolerance; + } + + public void refresh() { + this.config = log.config(); + } + + +} diff --git a/core/src/main/java/kafka/automq/table/worker/Writer.java b/core/src/main/java/kafka/automq/table/worker/Writer.java new file mode 100644 index 0000000000..8256799812 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/Writer.java @@ -0,0 +1,122 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.events.PartitionMetric; +import kafka.automq.table.events.TopicMetric; + +import org.apache.iceberg.io.WriteResult; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; + +public interface Writer { + + /** + * Write record to memory + */ + void write(int partition, org.apache.kafka.common.record.Record kafkaRecord) throws IOException; + + /** + * Asynchronously flush the data from memory to persistence storage + */ + CompletableFuture flush(FlushMode flushMode, ExecutorService flushExecutor, Executor eventLoop); + + /** + * Abort the writer and clean up the inflight resources. + */ + void abort() throws IOException; + + /** + * Complete the writer and return the results + */ + List complete() throws IOException; + + /** + * Get the results of the completed writer + */ + List results(); + + /** + * Check if the writer is completed + */ + boolean isCompleted(); + + /** + * Check if the writer is full if full should switch to a new writer. + */ + boolean isFull(); + + /** + * Get partition to the offset range map . + */ + Map getOffsets(); + + /** + * Get partition offset range. + */ + IcebergWriter.OffsetRange getOffset(int partition); + + /** + * Set partition initial offset range [offset, offset). + */ + void setOffset(int partition, long offset); + + /** + * Set new end offset for the partition. + */ + void setEndOffset(int partition, long offset); + + /** + * Get in memory dirty bytes. + */ + long dirtyBytes(); + + void updateWatermark(int partition, long timestamp); + + TopicMetric topicMetric(); + + Map partitionMetrics(); + + int targetFileSize(); + + class OffsetRange { + long start; + long end; + + public OffsetRange(long start) { + this.start = start; + this.end = start; + } + + public long start() { + return start; + } + + public long end() { + return end; + } + + } +} diff --git a/core/src/main/java/kafka/automq/table/worker/WriterFactory.java b/core/src/main/java/kafka/automq/table/worker/WriterFactory.java new file mode 100644 index 0000000000..acacf54d07 --- /dev/null +++ b/core/src/main/java/kafka/automq/table/worker/WriterFactory.java @@ -0,0 +1,31 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import org.apache.iceberg.PartitionSpec; + +public interface WriterFactory { + + Writer newWriter(); + + PartitionSpec partitionSpec(); + + void reset(); +} diff --git a/core/src/main/java/kafka/automq/utils/ClientUtils.java b/core/src/main/java/kafka/automq/utils/ClientUtils.java new file mode 100644 index 0000000000..0bd3d32a2a --- /dev/null +++ b/core/src/main/java/kafka/automq/utils/ClientUtils.java @@ -0,0 +1,87 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.utils; + +import kafka.cluster.EndPoint; +import kafka.server.KafkaConfig; + +import org.apache.kafka.common.network.ListenerName; +import org.apache.kafka.common.security.auth.SecurityProtocol; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Properties; + +import static scala.jdk.javaapi.CollectionConverters.asJava; + +public class ClientUtils { + public static Properties clusterClientBaseConfig(KafkaConfig kafkaConfig) { + ListenerName listenerName = kafkaConfig.interBrokerListenerName(); + + List endpoints = asJava(kafkaConfig.effectiveAdvertisedBrokerListeners()); + Optional endpointOpt = endpoints.stream().filter(e -> listenerName.equals(e.listenerName())).findFirst(); + if (endpointOpt.isEmpty()) { + throw new IllegalArgumentException("Cannot find " + listenerName + " in endpoints " + endpoints); + } + + EndPoint endpoint = endpointOpt.get(); + SecurityProtocol securityProtocol = kafkaConfig.interBrokerSecurityProtocol(); + Map parsedConfigs = kafkaConfig.valuesWithPrefixOverride(listenerName.configPrefix()); + + // mirror ChannelBuilders#channelBuilderConfigs + kafkaConfig.originals().entrySet().stream() + .filter(entry -> !parsedConfigs.containsKey(entry.getKey())) + // exclude already parsed listener prefix configs + .filter(entry -> !(entry.getKey().startsWith(listenerName.configPrefix()) + && parsedConfigs.containsKey(entry.getKey().substring(listenerName.configPrefix().length())))) + // exclude keys like `{mechanism}.some.prop` if "listener.name." prefix is present and key `some.prop` exists in parsed configs. + .filter(entry -> !parsedConfigs.containsKey(entry.getKey().substring(entry.getKey().indexOf('.') + 1))) + .forEach(entry -> parsedConfigs.put(entry.getKey(), entry.getValue())); + + Properties clientConfig = new Properties(); + parsedConfigs.entrySet().stream() + .filter(entry -> entry.getValue() != null) + .filter(entry -> isSecurityKey(entry.getKey(), listenerName)) + .forEach(entry -> clientConfig.put(entry.getKey(), entry.getValue())); + + String interBrokerSaslMechanism = kafkaConfig.saslMechanismInterBrokerProtocol(); + if (interBrokerSaslMechanism != null && !interBrokerSaslMechanism.isEmpty()) { + kafkaConfig.originalsWithPrefix(listenerName.saslMechanismConfigPrefix(interBrokerSaslMechanism)).entrySet().stream() + .filter(entry -> entry.getValue() != null) + .forEach(entry -> clientConfig.put(entry.getKey(), entry.getValue())); + clientConfig.putIfAbsent("sasl.mechanism", interBrokerSaslMechanism); + } + + clientConfig.put("security.protocol", securityProtocol.toString()); + clientConfig.put("bootstrap.servers", String.format("%s:%d", endpoint.host(), endpoint.port())); + return clientConfig; + } + + // Filter out non-security broker options (e.g. compression.type, log.retention.hours) so internal clients + // only inherit listener-specific SSL/SASL settings. + private static boolean isSecurityKey(String key, ListenerName listenerName) { + return key.startsWith("ssl.") + || key.startsWith("sasl.") + || key.startsWith("security.") + || key.startsWith(listenerName.configPrefix()); + } + +} diff --git a/core/src/main/java/kafka/automq/utils/JsonUtils.java b/core/src/main/java/kafka/automq/utils/JsonUtils.java new file mode 100644 index 0000000000..5c814a16c2 --- /dev/null +++ b/core/src/main/java/kafka/automq/utils/JsonUtils.java @@ -0,0 +1,65 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.utils; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; + +import java.util.Optional; + +public class JsonUtils { + private static final ObjectMapper MAPPER = new ObjectMapper(); + + public static String toArray(String s) { + ArrayNode array = MAPPER.createArrayNode(); + array.add(s); + try { + return MAPPER.writeValueAsString(array); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException(String.format("Failed to convert %s to json array", s), e); + } + } + + public static String encode(Object obj) { + try { + return MAPPER.writeValueAsString(obj); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException(e); + } + } + + public static T decode(String raw, Class clazz) { + try { + return MAPPER.readValue(raw, clazz); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("json parse (" + raw + ") fail", e); + } + } + + public static Optional getValue(String json, String key) { + try { + return Optional.ofNullable(MAPPER.readTree(json).get(key)).map(JsonNode::asText); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("json parse (" + json + ") fail", e); + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/AsyncSender.java b/core/src/main/java/kafka/automq/zerozone/AsyncSender.java new file mode 100644 index 0000000000..f6775d3620 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/AsyncSender.java @@ -0,0 +1,230 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.server.KafkaConfig; + +import org.apache.kafka.clients.ApiVersions; +import org.apache.kafka.clients.ClientRequest; +import org.apache.kafka.clients.ClientResponse; +import org.apache.kafka.clients.ManualMetadataUpdater; +import org.apache.kafka.clients.MetadataRecoveryStrategy; +import org.apache.kafka.clients.NetworkClient; +import org.apache.kafka.clients.NetworkClientUtils; +import org.apache.kafka.clients.RequestCompletionHandler; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.network.ChannelBuilder; +import org.apache.kafka.common.network.ChannelBuilders; +import org.apache.kafka.common.network.NetworkReceive; +import org.apache.kafka.common.network.Selectable; +import org.apache.kafka.common.network.Selector; +import org.apache.kafka.common.requests.AbstractRequest; +import org.apache.kafka.common.security.JaasContext; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.common.utils.Time; + +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.SocketTimeoutException; +import java.util.Collections; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; + +public interface AsyncSender { + + CompletableFuture sendRequest( + Node node, + AbstractRequest.Builder requestBuilder + ); + + void initiateClose(); + + void close(); + + class BrokersAsyncSender implements AsyncSender { + private static final Logger LOGGER = LoggerFactory.getLogger(BrokersAsyncSender.class); + private final NetworkClient networkClient; + private final Time time; + private final ExecutorService executorService; + private final AtomicBoolean shouldRun = new AtomicBoolean(true); + + public BrokersAsyncSender( + KafkaConfig brokerConfig, + Metrics metrics, + String metricGroupPrefix, + Time time, + String clientId, + LogContext logContext + ) { + + ChannelBuilder channelBuilder = ChannelBuilders.clientChannelBuilder( + brokerConfig.interBrokerSecurityProtocol(), + JaasContext.Type.SERVER, + brokerConfig, + brokerConfig.interBrokerListenerName(), + brokerConfig.saslMechanismInterBrokerProtocol(), + time, + brokerConfig.saslInterBrokerHandshakeRequestEnable(), + logContext + ); + Selector selector = new Selector( + NetworkReceive.UNLIMITED, + brokerConfig.connectionsMaxIdleMs(), + metrics, + time, + metricGroupPrefix, + Collections.emptyMap(), + false, + channelBuilder, + logContext + ); + this.networkClient = new NetworkClient( + selector, + new ManualMetadataUpdater(), + clientId, + 64, + 0, + 0, + Selectable.USE_DEFAULT_BUFFER_SIZE, + brokerConfig.replicaSocketReceiveBufferBytes(), + brokerConfig.requestTimeoutMs(), + brokerConfig.connectionSetupTimeoutMs(), + brokerConfig.connectionSetupTimeoutMaxMs(), + time, + true, + new ApiVersions(), + logContext, + MetadataRecoveryStrategy.REBOOTSTRAP + ); + this.time = time; + executorService = Threads.newFixedThreadPoolWithMonitor(1, metricGroupPrefix, true, LOGGER); + executorService.submit(this::run); + } + + private final ConcurrentMap> waitingSendRequests = new ConcurrentHashMap<>(); + + @Override + public CompletableFuture sendRequest(Node node, + AbstractRequest.Builder requestBuilder) { + CompletableFuture cf = new CompletableFuture<>(); + waitingSendRequests.compute(node, (n, queue) -> { + if (queue == null) { + queue = new ConcurrentLinkedQueue<>(); + } + queue.add(new Request(requestBuilder, cf)); + return queue; + }); + return cf; + } + + private void run() { + Map connectingStates = new ConcurrentHashMap<>(); + while (shouldRun.get()) { + // TODO: graceful shutdown + try { + long now = time.milliseconds(); + waitingSendRequests.forEach((node, queue) -> { + if (queue.isEmpty()) { + return; + } + if (NetworkClientUtils.isReady(networkClient, node, now)) { + connectingStates.remove(node); + Request request = queue.poll(); + ClientRequest clientRequest = networkClient.newClientRequest(Integer.toString(node.id()), request.requestBuilder, now, true, 10000, new RequestCompletionHandler() { + @Override + public void onComplete(ClientResponse response) { + request.cf.complete(response); + } + }); + networkClient.send(clientRequest, now); + } else { + ConnectingState connectingState = connectingStates.get(node); + if (connectingState == null) { + networkClient.ready(node, now); + connectingStates.put(node, new ConnectingState(now)); + } else { + if (now - connectingState.startConnectNanos > 3000) { + for (; ; ) { + Request request = queue.poll(); + if (request == null) { + break; + } + request.cf.completeExceptionally(new SocketTimeoutException(String.format("Cannot connect to node=%s", node))); + } + connectingStates.remove(node); + } else if (now - connectingState.startConnectNanos > 1000 && connectingState.connectTimes < 2) { + // The broker network maybe slightly ready after the broker become UNFENCED. + // So we need to retry connect twice. + networkClient.ready(node, now); + connectingState.connectTimes = connectingState.connectTimes + 1; + } + } + } + }); + networkClient.poll(1, now); + } catch (Throwable e) { + LOGGER.error("Processor get uncaught exception", e); + } + } + } + + @Override + public void initiateClose() { + networkClient.initiateClose(); + } + + @Override + public void close() { + networkClient.close(); + shouldRun.set(false); + } + } + + class Request { + final AbstractRequest.Builder requestBuilder; + final CompletableFuture cf; + + public Request(AbstractRequest.Builder requestBuilder, CompletableFuture cf) { + this.requestBuilder = requestBuilder; + this.cf = cf; + } + } + + class ConnectingState { + final long startConnectNanos; + int connectTimes; + + public ConnectingState(long startConnectNanos) { + this.startConnectNanos = startConnectNanos; + connectTimes = 1; + } + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/ChannelOffset.java b/core/src/main/java/kafka/automq/zerozone/ChannelOffset.java new file mode 100644 index 0000000000..a31a40d1d9 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ChannelOffset.java @@ -0,0 +1,90 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class ChannelOffset { + private static final byte MAGIC = (byte) 0x86; + private static final int CHANNEL_ID_INDEX = 1; + private static final int ORDER_HINT_INDEX = 3; + private static final int CHANNEL_OWNER_NODE_ID_INDEX = 5; + private static final int CHANNEL_ATTRIBUTES = 9; + private static final int WAL_RECORD_OFFSET_INDEX = 13; + + private final ByteBuf buf; + + private ChannelOffset(ByteBuf buf) { + this.buf = buf; + } + + public static ChannelOffset of(ByteBuf buf) { + return new ChannelOffset(buf); + } + + public static ChannelOffset of(short channelId, short orderHint, int channelOwnerNodeId, int attributes, + ByteBuf walRecordOffset) { + ByteBuf channelOffset = Unpooled.buffer(1 /* magic */ + 2 /* channelId */ + 2 /* orderHint */ + + 4 /* channelOwnerNodeId */ + 4 /* attributes */ + walRecordOffset.readableBytes()); + channelOffset.writeByte(MAGIC); + channelOffset.writeShort(channelId); + channelOffset.writeShort(orderHint); + channelOffset.writeInt(channelOwnerNodeId); + channelOffset.writeInt(attributes); + channelOffset.writeBytes(walRecordOffset.duplicate()); + return of(channelOffset); + } + + public short channelId() { + return buf.getShort(CHANNEL_ID_INDEX); + } + + public short orderHint() { + return buf.getShort(ORDER_HINT_INDEX); + } + + public int channelOwnerNodeId() { + return buf.getInt(CHANNEL_OWNER_NODE_ID_INDEX); + } + + public int attributes() { + return buf.getInt(CHANNEL_ATTRIBUTES); + } + + public ByteBuf walRecordOffset() { + return buf.slice(WAL_RECORD_OFFSET_INDEX, buf.readableBytes() - WAL_RECORD_OFFSET_INDEX); + } + + public ByteBuf byteBuf() { + return buf; + } + + @Override + public String toString() { + return "ChannelOffset{" + + "channelId=" + channelId() + + ", orderHint=" + orderHint() + + ", channelOwnerNodeId=" + channelOwnerNodeId() + + ", attributes=" + attributes() + + ", walRecordOffset=" + walRecordOffset() + + '}'; + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ClientRackProvider.java b/core/src/main/java/kafka/automq/zerozone/ClientRackProvider.java new file mode 100644 index 0000000000..735dc7f8ea --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ClientRackProvider.java @@ -0,0 +1,28 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; + +public interface ClientRackProvider { + + String rack(ClientIdMetadata clientId); + +} diff --git a/core/src/main/java/kafka/automq/zerozone/CommittedEpochManager.java b/core/src/main/java/kafka/automq/zerozone/CommittedEpochManager.java new file mode 100644 index 0000000000..56a68e7043 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/CommittedEpochManager.java @@ -0,0 +1,145 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.controller.stream.NodeCommittedEpoch; +import org.apache.kafka.controller.stream.RouterChannelEpoch; + +import com.automq.stream.Context; +import com.automq.stream.api.KeyValue; +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Iterator; +import java.util.Map; +import java.util.NavigableMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +public class CommittedEpochManager implements RouterChannelProvider.EpochListener { + private static final Logger LOGGER = LoggerFactory.getLogger(CommittedEpochManager.class); + + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock(); + private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock(); + private RouterChannelEpoch routerChannelEpoch; + private long waitingCommitEpoch = -1L; + private long committedEpoch = -1L; + private CompletableFuture commitCf = CompletableFuture.completedFuture(null); + private final NavigableMap epoch2inflight = new ConcurrentSkipListMap<>(); + + private volatile ScheduledFuture tryBumpCommitedEpochFuture; + + private final int nodeId; + + public CommittedEpochManager(int nodeId) { + this.nodeId = nodeId; + } + + public ReentrantReadWriteLock.ReadLock readLock() { + return readLock; + } + + public AtomicLong epochInflight(long epoch) { + return epoch2inflight.computeIfAbsent(epoch, e -> new AtomicLong()); + } + + private void start() { + this.tryBumpCommitedEpochFuture = Threads.COMMON_SCHEDULER.scheduleWithFixedDelay(this::tryBumpCommittedEpoch, 1, 1, TimeUnit.SECONDS); + } + + public void close() { + if (tryBumpCommitedEpochFuture != null) { + tryBumpCommitedEpochFuture.cancel(true); + } + } + + private void tryBumpCommittedEpoch() { + writeLock.lock(); + try { + tryBumpCommittedEpoch0(); + } catch (Exception e) { + LOGGER.error("Error while trying bumping committed epoch", e); + } finally { + writeLock.unlock(); + } + } + + private void tryBumpCommittedEpoch0() { + long fencedEpoch = routerChannelEpoch.getFenced(); + Iterator> it = epoch2inflight.entrySet().iterator(); + long newWaitingEpoch = waitingCommitEpoch; + while (it.hasNext()) { + Map.Entry entry = it.next(); + long epoch = entry.getKey(); + if (epoch > fencedEpoch) { + break; + } + AtomicLong inflight = entry.getValue(); + if (inflight.get() <= 0) { + // We only bump the commitEpoch when this epoch was fenced and has no inflight requests. + it.remove(); + newWaitingEpoch = epoch; + } else { + break; + } + } + if (epoch2inflight.isEmpty()) { + newWaitingEpoch = fencedEpoch; + } + if (newWaitingEpoch != waitingCommitEpoch) { + this.waitingCommitEpoch = newWaitingEpoch; + } + if (commitCf.isDone() && waitingCommitEpoch != committedEpoch) { + long newCommittedEpoch = waitingCommitEpoch; + commitCf = commitCf + .thenCompose(nil -> Context.instance().confirmWAL().commit(TimeUnit.SECONDS.toMillis(10))) + .thenCompose(nil -> + Context.instance().kvClient().putKV(KeyValue.of( + NodeCommittedEpoch.NODE_COMMITED_EPOCH_KEY_PREFIX + nodeId, + NodeCommittedEpoch.encode(new NodeCommittedEpoch(newCommittedEpoch), (short) 0).nioBuffer() + )).thenAccept(rst -> committedEpoch = newCommittedEpoch) + ).exceptionally(ex -> { + LOGGER.error("[BUMP_COMMITTED_EPOCH_FAIL]", ex); + return null; + }); + } + } + + @Override + public void onNewEpoch(RouterChannelEpoch epoch) { + writeLock.lock(); + try { + boolean first = routerChannelEpoch == null; + routerChannelEpoch = epoch; + if (first) { + start(); + } + } finally { + writeLock.unlock(); + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ConfirmWALProvider.java b/core/src/main/java/kafka/automq/zerozone/ConfirmWALProvider.java new file mode 100644 index 0000000000..2c78eb098e --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ConfirmWALProvider.java @@ -0,0 +1,30 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.s3.wal.WriteAheadLog; + +import java.util.concurrent.CompletableFuture; + +public interface ConfirmWALProvider { + + CompletableFuture readOnly(String walConfig, int nodeId); + +} diff --git a/core/src/main/java/kafka/automq/zerozone/DefaultClientRackProvider.java b/core/src/main/java/kafka/automq/zerozone/DefaultClientRackProvider.java new file mode 100644 index 0000000000..d7e535bbf7 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/DefaultClientRackProvider.java @@ -0,0 +1,181 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.server.DynamicBrokerConfig; + +import org.apache.kafka.common.Reconfigurable; +import org.apache.kafka.common.config.ConfigException; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +public class DefaultClientRackProvider implements ClientRackProvider, Reconfigurable { + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultClientRackProvider.class); + private static final String ZONE_CIDR_BLOCKS_CONFIG_KEY = "automq.zone.cidr.blocks"; + private static final Set RECONFIGURABLE_CONFIGS; + private CIDRMatcher cidrMatcher = new CIDRMatcher(""); + + static { + RECONFIGURABLE_CONFIGS = Set.of( + ZONE_CIDR_BLOCKS_CONFIG_KEY + ); + RECONFIGURABLE_CONFIGS.forEach(DynamicBrokerConfig.AllDynamicConfigs()::add); + } + + @Override + public String rack(ClientIdMetadata clientId) { + String rack = clientId.rack(); + if (rack != null) { + return rack; + } + CIDRBlock block = cidrMatcher.find(clientId.clientAddress().getHostAddress()); + if (block == null) { + return null; + } + return block.zone(); + } + + @Override + public Set reconfigurableConfigs() { + return Set.of(ZONE_CIDR_BLOCKS_CONFIG_KEY); + } + + @Override + public void validateReconfiguration(Map map) throws ConfigException { + config(map, true); + } + + @Override + public void reconfigure(Map map) { + config(map, false); + } + + @Override + public void configure(Map map) { + config(map, false); + } + + private void config(Map map, boolean validate) { + String zoneCidrBlocksConfig = (String) map.get(ZONE_CIDR_BLOCKS_CONFIG_KEY); + if (zoneCidrBlocksConfig != null) { + CIDRMatcher matcher = new CIDRMatcher(zoneCidrBlocksConfig); + if (!validate) { + cidrMatcher = matcher; + LOGGER.info("apply new zone CIDR blocks {}", zoneCidrBlocksConfig); + } + } + } + + public static class CIDRMatcher { + private final Map> maskLength2blocks = new HashMap<>(); + private final List reverseMaskLengthList = new ArrayList<>(); + + public CIDRMatcher(String config) { + for (String cidrBlocksOfZone : config.split("<>")) { + String[] parts = cidrBlocksOfZone.split("@"); + if (parts.length != 2) { + continue; + } + String zone = parts[0]; + String[] cidrList = parts[1].split(","); + for (String cidr : cidrList) { + CIDRBlock block = parseCidr(cidr, zone); + maskLength2blocks + .computeIfAbsent(block.prefixLength, k -> new ArrayList<>()) + .add(block); + } + } + reverseMaskLengthList.addAll(maskLength2blocks.keySet()); + reverseMaskLengthList.sort(Comparator.reverseOrder()); + } + + public CIDRBlock find(String ip) { + long ipLong = ipToLong(ip); + for (int prefix : reverseMaskLengthList) { + List blocks = maskLength2blocks.get(prefix); + if (blocks != null) { + for (CIDRBlock block : blocks) { + if (block.contains(ipLong)) { + return block; + } + } + } + } + return null; + } + + private CIDRBlock parseCidr(String cidr, String zone) { + String[] parts = cidr.split("/"); + String ip = parts[0]; + int maskLength = Integer.parseInt(parts[1]); + long ipLong = ipToLong(ip); + long mask = (0xFFFFFFFFL << (32 - maskLength)) & 0xFFFFFFFFL; + long networkAddress = ipLong & mask; + return new CIDRBlock(cidr, networkAddress, mask, maskLength, zone); + } + + private long ipToLong(String ipAddress) { + String[] octets = ipAddress.split("\\."); + long result = 0; + for (String octet : octets) { + result = (result << 8) | Integer.parseUnsignedInt(octet); + } + return result; + } + + } + + public static class CIDRBlock { + private final String cidr; + private final long networkAddress; + private final long mask; + final int prefixLength; + private final String zone; + + public CIDRBlock(String cidr, long networkAddress, long mask, int prefixLength, String zone) { + this.cidr = cidr; + this.networkAddress = networkAddress; + this.mask = mask; + this.prefixLength = prefixLength; + this.zone = zone; + } + + public boolean contains(long ip) { + return (ip & mask) == networkAddress; + } + + public String cidr() { + return cidr; + } + + public String zone() { + return zone; + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/DefaultConfirmWALProvider.java b/core/src/main/java/kafka/automq/zerozone/DefaultConfirmWALProvider.java new file mode 100644 index 0000000000..b1e6d9c778 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/DefaultConfirmWALProvider.java @@ -0,0 +1,79 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; +import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.s3.wal.WriteAheadLog; +import com.automq.stream.s3.wal.impl.object.ObjectWALConfig; +import com.automq.stream.s3.wal.impl.object.ObjectWALService; +import com.automq.stream.utils.Time; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; + +public class DefaultConfirmWALProvider implements ConfirmWALProvider { + private static final ExecutorService START_EXECUTOR = java.util.concurrent.Executors.newCachedThreadPool(); + private final Map objectStorages = new ConcurrentHashMap<>(); + private final String clusterId; + private final Time time = Time.SYSTEM; + + public DefaultConfirmWALProvider(String clusterId) { + this.clusterId = clusterId; + } + + @Override + public CompletableFuture readOnly(String walConfig, int nodeId) { + BucketURI bucketURI = BucketURI.parse(walConfig); + ObjectStorage objectStorage = objectStorages.computeIfAbsent(bucketURI.bucketId(), id -> { + try { + return ObjectStorageFactory.instance().builder(bucketURI).readWriteIsolate(false) + .inboundLimiter(GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.INBOUND)) + .outboundLimiter(GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.OUTBOUND)) + .build(); + } catch (IllegalArgumentException e) { + return null; + } + } + ); + if (objectStorage == null) { + throw new IllegalArgumentException("Cannot parse " + walConfig); + } + ObjectWALConfig objectWALConfig = ObjectWALConfig.builder() + .withClusterId(clusterId) + .withNodeId(nodeId) + .withOpenMode(OpenMode.READ_ONLY) + .build(); + ObjectWALService wal = new ObjectWALService(time, objectStorage, objectWALConfig); + return CompletableFuture.runAsync(() -> { + try { + wal.start(); + } catch (Throwable e) { + throw new RuntimeException(e); + } + }, START_EXECUTOR).thenApply(v -> wal); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/DefaultLinkRecordDecoder.java b/core/src/main/java/kafka/automq/zerozone/DefaultLinkRecordDecoder.java new file mode 100644 index 0000000000..004bb2cf68 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/DefaultLinkRecordDecoder.java @@ -0,0 +1,78 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.MutableRecordBatch; + +import com.automq.stream.s3.cache.SnapshotReadCache; +import com.automq.stream.s3.model.StreamRecordBatch; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.CompletableFuture; + +import io.netty.buffer.ByteBuf; + +public class DefaultLinkRecordDecoder implements com.automq.stream.api.LinkRecordDecoder { + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultLinkRecordDecoder.class); + private final RouterChannelProvider channelProvider; + + public DefaultLinkRecordDecoder(RouterChannelProvider channelProvider) { + this.channelProvider = channelProvider; + } + + @Override + public int decodedSize(ByteBuf linkRecordBuf) { + return LinkRecord.decodedSize(linkRecordBuf); + } + + @Override + public CompletableFuture decode(StreamRecordBatch src) { + try { + LinkRecord linkRecord = LinkRecord.decode(src.getPayload()); + ChannelOffset channelOffset = linkRecord.channelOffset(); + RouterChannel routerChannel = channelProvider.readOnlyChannel(channelOffset.channelOwnerNodeId()); + return routerChannel.get(channelOffset.byteBuf()).thenApply(buf -> { + try (ZoneRouterProduceRequest req = ZoneRouterPackReader.decodeDataBlock(buf).get(0)) { + MemoryRecords records = (MemoryRecords) (req.data().topicData().iterator().next() + .partitionData().iterator().next() + .records()); + MutableRecordBatch recordBatch = records.batches().iterator().next(); + recordBatch.setLastOffset(linkRecord.lastOffset()); + recordBatch.setMaxTimestamp(linkRecord.timestampType(), linkRecord.maxTimestamp()); + recordBatch.setPartitionLeaderEpoch(linkRecord.partitionLeaderEpoch()); + return StreamRecordBatch.of(src.getStreamId(), src.getEpoch(), src.getBaseOffset(), + -src.getCount(), records.buffer(), SnapshotReadCache.ENCODE_ALLOC); + } finally { + buf.release(); + } + }).whenComplete((rst, ex) -> { + src.release(); + if (ex != null) { + LOGGER.error("Error while decoding link record, link={}", linkRecord, ex); + } + }); + } catch (Throwable t) { + return CompletableFuture.failedFuture(t); + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/DefaultReplayer.java b/core/src/main/java/kafka/automq/zerozone/DefaultReplayer.java new file mode 100644 index 0000000000..dedfefea78 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/DefaultReplayer.java @@ -0,0 +1,62 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.Context; +import com.automq.stream.s3.cache.SnapshotReadCache; +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.WriteAheadLog; + +import java.util.List; +import java.util.concurrent.CompletableFuture; + +public class DefaultReplayer implements Replayer { + private SnapshotReadCache.EventListener listener; + private SnapshotReadCache snapshotReadCache; + + @Override + public CompletableFuture replay(List objects) { + return snapshotReadCache().replay(objects); + } + + @Override + public CompletableFuture replay(WriteAheadLog confirmWAL, RecordOffset startOffset, RecordOffset endOffset, List walRecords) { + return snapshotReadCache().replay(confirmWAL, startOffset, endOffset, walRecords); + } + + private SnapshotReadCache snapshotReadCache() { + if (snapshotReadCache == null) { + snapshotReadCache = Context.instance().snapshotReadCache(); + if (listener != null) { + snapshotReadCache.addEventListener(listener); + } + } + return snapshotReadCache; + } + + public void setCacheEventListener(SnapshotReadCache.EventListener listener) { + this.listener = listener; + if (snapshotReadCache != null) { + snapshotReadCache.addEventListener(listener); + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java b/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java new file mode 100644 index 0000000000..31ad44d8ba --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java @@ -0,0 +1,164 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.controller.stream.RouterChannelEpoch; +import org.apache.kafka.image.MetadataDelta; +import org.apache.kafka.image.MetadataImage; + +import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; +import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.s3.wal.impl.object.ObjectWALConfig; +import com.automq.stream.s3.wal.impl.object.ObjectWALService; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Time; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; + +import io.netty.buffer.Unpooled; + +public class DefaultRouterChannelProvider implements RouterChannelProvider { + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultRouterChannelProvider.class); + public static final String WAL_TYPE = "rc"; + private final int nodeId; + private final long nodeEpoch; + private final short channelId; + private final BucketURI bucketURI; + private volatile RouterChannel routerChannel; + private ObjectStorage objectStorage; + private final Map routerChannels = new ConcurrentHashMap<>(); + private final String clusterId; + + private final List epochListeners = new CopyOnWriteArrayList<>(); + private volatile RouterChannelEpoch epoch = new RouterChannelEpoch(-3L, -2L, 0, 0); + + public DefaultRouterChannelProvider(int nodeId, long nodeEpoch, BucketURI bucketURI, String clusterId) { + this.nodeId = nodeId; + this.nodeEpoch = nodeEpoch; + this.bucketURI = bucketURI; + this.channelId = bucketURI.bucketId(); + this.clusterId = clusterId; + } + + @Override + public RouterChannel channel() { + if (routerChannel != null) { + return routerChannel; + } + synchronized (this) { + if (routerChannel == null) { + ObjectWALConfig config = ObjectWALConfig.builder() + .withClusterId(clusterId) + .withNodeId(nodeId) + .withEpoch(nodeEpoch) + .withOpenMode(OpenMode.READ_WRITE) + .withType(WAL_TYPE) + .build(); + ObjectWALService wal = new ObjectWALService(Time.SYSTEM, objectStorage(), config); + RouterChannel routerChannel = new ObjectRouterChannel(this.nodeId, channelId, wal); + routerChannel.nextEpoch(epoch.getCurrent()); + routerChannel.trim(epoch.getCommitted()); + this.routerChannel = routerChannel; + } + return routerChannel; + } + } + + @Override + public RouterChannel readOnlyChannel(int node) { + if (nodeId == node) { + return channel(); + } + return routerChannels.computeIfAbsent(node, nodeId -> { + ObjectWALConfig config = ObjectWALConfig.builder().withClusterId(clusterId).withNodeId(node).withOpenMode(OpenMode.READ_ONLY).withType(WAL_TYPE).build(); + ObjectWALService wal = new ObjectWALService(Time.SYSTEM, objectStorage(), config); + return new ObjectRouterChannel(nodeId, channelId, wal); + }); + } + + @Override + public RouterChannelEpoch epoch() { + return epoch; + } + + @Override + public void addEpochListener(EpochListener listener) { + epochListeners.add(listener); + } + + @Override + public void close() { + FutureUtil.suppress(() -> routerChannel.close().get(), LOGGER); + routerChannels.forEach((nodeId, channel) -> FutureUtil.suppress(() -> channel.close().get(), LOGGER)); + } + + @Override + public void onChange(MetadataDelta delta, MetadataImage image) { + if (delta.kvDelta() == null) { + return; + } + ByteBuffer value = delta.kvDelta().changedKV().get(RouterChannelEpoch.ROUTER_CHANNEL_EPOCH_KEY); + if (value == null) { + return; + } + synchronized (this) { + this.epoch = RouterChannelEpoch.decode(Unpooled.wrappedBuffer(value.slice())); + RouterChannel routerChannel = this.routerChannel; + if (routerChannel != null) { + routerChannel.nextEpoch(epoch.getCurrent()); + routerChannel.trim(epoch.getCommitted()); + } + } + notifyEpochListeners(epoch); + + } + + private void notifyEpochListeners(RouterChannelEpoch epoch) { + for (EpochListener listener : epochListeners) { + try { + listener.onNewEpoch(epoch); + } catch (Throwable t) { + LOGGER.error("Failed to notify epoch listener {}", listener, t); + } + } + } + + synchronized ObjectStorage objectStorage() { + if (objectStorage == null) { + this.objectStorage = ObjectStorageFactory.instance().builder(bucketURI) + .readWriteIsolate(true) + .inboundLimiter(GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.INBOUND)) + .outboundLimiter(GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.OUTBOUND)) + .build(); + } + return objectStorage; + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/GetRouterOutNode.java b/core/src/main/java/kafka/automq/zerozone/GetRouterOutNode.java new file mode 100644 index 0000000000..afa4a8881f --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/GetRouterOutNode.java @@ -0,0 +1,28 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; + +import org.apache.kafka.common.Node; + +public interface GetRouterOutNode { + Node getRouteOutNode(String topicName, int partition, ClientIdMetadata clientId); +} diff --git a/core/src/main/java/kafka/automq/zerozone/LinkRecord.java b/core/src/main/java/kafka/automq/zerozone/LinkRecord.java new file mode 100644 index 0000000000..7f20509e5b --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/LinkRecord.java @@ -0,0 +1,131 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.MutableRecordBatch; +import org.apache.kafka.common.record.TimestampType; + +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class LinkRecord { + private static final byte MAGIC_V0 = (byte) 0x00; + private static final int CHANNEL_OFFSET_OFFSET = 1 /* magic */ + 8 /* last offset */ + 4 /* timestamp type */ + 8 /* max timestamp */ + 4 /* leader epoch */; + private final long lastOffset; + private final TimestampType timestampType; + private final long maxTimestamp; + private final int partitionLeaderEpoch; + private final ChannelOffset channelOffset; + + public LinkRecord(long lastOffset, TimestampType timestampType, long maxTimestamp, int partitionLeaderEpoch, + ChannelOffset channelOffset) { + this.lastOffset = lastOffset; + this.timestampType = timestampType; + this.maxTimestamp = maxTimestamp; + this.partitionLeaderEpoch = partitionLeaderEpoch; + this.channelOffset = channelOffset; + } + + public long lastOffset() { + return lastOffset; + } + + public TimestampType timestampType() { + return timestampType; + } + + public long maxTimestamp() { + return maxTimestamp; + } + + public int partitionLeaderEpoch() { + return partitionLeaderEpoch; + } + + public ChannelOffset channelOffset() { + return channelOffset; + } + + @Override + public String toString() { + return "LinkRecord{" + + "lastOffset=" + lastOffset + + ", timestampType=" + timestampType + + ", maxTimestamp=" + maxTimestamp + + ", partitionLeaderEpoch=" + partitionLeaderEpoch + + ", channelOffset=" + channelOffset + + '}'; + } + + public static ByteBuf encode(ChannelOffset channelOffset, MemoryRecords memoryRecords) { + // The MemoryRecords only contains one RecordBatch, cause of produce only send one RecordBatch per partition. + if (channelOffset == null) { + return null; + } + MutableRecordBatch recordBatch = memoryRecords.batches().iterator().next(); + long offset = recordBatch.lastOffset(); + long timestamp = recordBatch.maxTimestamp(); + int partitionLeaderEpoch = recordBatch.partitionLeaderEpoch(); + + ByteBuf buffer = Unpooled.buffer(1 /* magic */ + 8 /* lastOffset */ + 4 /* timestampType */ + 8 /* maxTimestamp */ + + 4 /* partitionLeaderEpoch */ + channelOffset.byteBuf().readableBytes()); + buffer.writeByte(MAGIC_V0); + buffer.writeLong(offset); + buffer.writeInt(recordBatch.timestampType().id); + buffer.writeLong(timestamp); + buffer.writeInt(partitionLeaderEpoch); + buffer.writeBytes(channelOffset.byteBuf().slice()); + + return buffer; + } + + public static LinkRecord decode(ByteBuf buf) { + buf = buf.slice(); + byte magic = buf.readByte(); + if (magic != MAGIC_V0) { + throw new UnsupportedOperationException("Unsupported magic: " + magic); + } + long lastOffset = buf.readLong(); + TimestampType timestampType = TimestampType.forId(buf.readInt()); + long maxTimestamp = buf.readLong(); + int partitionLeaderEpoch = buf.readInt(); + ByteBuf channelOffset = Unpooled.buffer(buf.readableBytes()); + buf.readBytes(channelOffset); + return new LinkRecord(lastOffset, timestampType, maxTimestamp, partitionLeaderEpoch, ChannelOffset.of(channelOffset)); + } + + /** + * Get the size of the linked record. + */ + public static int decodedSize(ByteBuf linkRecordBuf) { + ByteBuf buf = linkRecordBuf.slice(); + byte magic = buf.getByte(0); + if (magic != MAGIC_V0) { + throw new UnsupportedOperationException("Unsupported magic: " + magic); + } + int channelOffsetSize = buf.readableBytes() - CHANNEL_OFFSET_OFFSET; + ByteBuf channelOffsetBuf = Unpooled.buffer(channelOffsetSize); + buf.getBytes(CHANNEL_OFFSET_OFFSET, channelOffsetBuf); + return DefaultRecordOffset.of(ChannelOffset.of(channelOffsetBuf).walRecordOffset()).size(); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/MismatchRecorder.java b/core/src/main/java/kafka/automq/zerozone/MismatchRecorder.java new file mode 100644 index 0000000000..f455bd67c5 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/MismatchRecorder.java @@ -0,0 +1,72 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; + +import org.apache.kafka.common.internals.Topic; + +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.TimeUnit; + +class MismatchRecorder { + private static final Logger LOGGER = LoggerFactory.getLogger(MismatchRecorder.class); + private static final int MAX_RECORD_SIZE = 16; + private ConcurrentMap topic2clientId = new ConcurrentHashMap<>(MAX_RECORD_SIZE); + private ConcurrentMap topic2clientId2 = new ConcurrentHashMap<>(MAX_RECORD_SIZE); + + private static final MismatchRecorder INSTANCE = new MismatchRecorder(); + + private MismatchRecorder() { + Threads.COMMON_SCHEDULER.scheduleWithFixedDelay(this::logAndReset, 1, 1, TimeUnit.MINUTES); + } + + public static MismatchRecorder instance() { + return INSTANCE; + } + + public void record(String topic, ClientIdMetadata clientId) { + if (topic2clientId.size() >= MAX_RECORD_SIZE) { + return; + } + if (topic.equals(Topic.AUTO_BALANCER_METRICS_TOPIC_NAME)) { + return; + } + topic2clientId.putIfAbsent(topic, clientId); + } + + private void logAndReset() { + if (topic2clientId.isEmpty()) { + return; + } + ConcurrentMap logMap = topic2clientId; + topic2clientId = topic2clientId2; + topic2clientId2 = logMap; + LOGGER.warn("[RACK_AWARE_MISMATCH],{}", logMap); + logMap.clear(); + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/NonBlockingLocalRouterHandler.java b/core/src/main/java/kafka/automq/zerozone/NonBlockingLocalRouterHandler.java new file mode 100644 index 0000000000..66da065605 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/NonBlockingLocalRouterHandler.java @@ -0,0 +1,33 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.message.AutomqZoneRouterResponseData; + +import java.util.concurrent.CompletableFuture; + +public interface NonBlockingLocalRouterHandler { + + CompletableFuture append( + ChannelOffset channelOffset, + ZoneRouterProduceRequest zoneRouterProduceRequest + ); + +} diff --git a/core/src/main/java/kafka/automq/zerozone/ObjectRouterChannel.java b/core/src/main/java/kafka/automq/zerozone/ObjectRouterChannel.java new file mode 100644 index 0000000000..42c957df27 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ObjectRouterChannel.java @@ -0,0 +1,168 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.trace.context.TraceContext; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; +import com.automq.stream.s3.wal.impl.object.ObjectWALService; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.LogContext; +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; + +import java.util.HashMap; +import java.util.LinkedList; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +import io.netty.buffer.ByteBuf; + +public class ObjectRouterChannel implements RouterChannel { + private static final ExecutorService ASYNC_EXECUTOR = Executors.newCachedThreadPool(); + private static final long OVER_CAPACITY_RETRY_DELAY_MS = 1000L; + private final Logger logger; + private final AtomicLong mockOffset = new AtomicLong(0); + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock(); + private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock(); + + private final ObjectWALService wal; + private final int nodeId; + private final short channelId; + + private long channelEpoch = 0L; + private final Queue channelEpochQueue = new LinkedList<>(); + private final Map channelEpoch2LastRecordOffset = new HashMap<>(); + + private final CompletableFuture startCf; + + public ObjectRouterChannel(int nodeId, short channelId, ObjectWALService wal) { + this.logger = new LogContext(String.format("[OBJECT_ROUTER_CHANNEL-%s-%s] ", channelId, nodeId)).logger(ObjectRouterChannel.class); + this.nodeId = nodeId; + this.channelId = channelId; + this.wal = wal; + this.startCf = CompletableFuture.runAsync(() -> { + try { + wal.start(); + } catch (Throwable e) { + logger.error("start object router channel failed.", e); + throw new RuntimeException(e); + } + }, ASYNC_EXECUTOR); + } + + @Override + public CompletableFuture append(int targetNodeId, short orderHint, ByteBuf data) { + return startCf.thenCompose(nil -> append0(targetNodeId, orderHint, data)); + } + + CompletableFuture append0(int targetNodeId, short orderHint, ByteBuf data) { + StreamRecordBatch record = StreamRecordBatch.of(targetNodeId, 0, mockOffset.incrementAndGet(), 1, data); + for (; ; ) { + try { + record.retain(); + return wal.append(TraceContext.DEFAULT, record).thenApply(walRst -> { + readLock.lock(); + try { + long epoch = this.channelEpoch; + ChannelOffset channelOffset = ChannelOffset.of(channelId, orderHint, nodeId, targetNodeId, walRst.recordOffset().buffer()); + channelEpoch2LastRecordOffset.put(epoch, walRst.recordOffset()); + return new AppendResult(epoch, channelOffset.byteBuf()); + } finally { + readLock.unlock(); + } + }).whenComplete((r, e) -> record.release()); + } catch (OverCapacityException e) { + logger.warn("OverCapacityException occurred while appending, err={}", e.getMessage()); + // Use block-based delayed retries for network backpressure. + Threads.sleep(OVER_CAPACITY_RETRY_DELAY_MS); + } catch (Throwable e) { + logger.error("[UNEXPECTED], append wal fail", e); + record.release(); + return CompletableFuture.failedFuture(e); + } + } + } + + @Override + public CompletableFuture get(ByteBuf channelOffset) { + return startCf.thenCompose(nil -> get0(channelOffset)); + } + + CompletableFuture get0(ByteBuf channelOffset) { + return wal.get(DefaultRecordOffset.of(ChannelOffset.of(channelOffset).walRecordOffset())).thenApply(streamRecordBatch -> { + ByteBuf payload = streamRecordBatch.getPayload().retainedSlice(); + streamRecordBatch.release(); + return payload; + }); + } + + @Override + public void nextEpoch(long epoch) { + writeLock.lock(); + try { + if (epoch > this.channelEpoch) { + this.channelEpochQueue.add(epoch); + this.channelEpoch = epoch; + } + } finally { + writeLock.unlock(); + } + } + + @Override + public void trim(long epoch) { + writeLock.lock(); + try { + RecordOffset recordOffset = null; + for (; ; ) { + Long channelEpoch = channelEpochQueue.peek(); + if (channelEpoch == null || channelEpoch > epoch) { + break; + } + channelEpochQueue.poll(); + RecordOffset removed = channelEpoch2LastRecordOffset.remove(channelEpoch); + if (removed != null) { + recordOffset = removed; + } + } + if (recordOffset != null) { + wal.trim(recordOffset); + logger.info("trim to epoch={} offset={}", epoch, recordOffset); + } + } finally { + writeLock.unlock(); + } + } + + @Override + public CompletableFuture close() { + return startCf.thenAcceptAsync(nil -> FutureUtil.suppress(wal::shutdownGracefully, logger), ASYNC_EXECUTOR); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/Position.java b/core/src/main/java/kafka/automq/zerozone/Position.java new file mode 100644 index 0000000000..118909d6b0 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/Position.java @@ -0,0 +1,38 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +public class Position { + private final int position; + private final int size; + + public Position(int position, int size) { + this.position = position; + this.size = size; + } + + public int position() { + return position; + } + + public int size() { + return size; + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ProxyNodeMapping.java b/core/src/main/java/kafka/automq/zerozone/ProxyNodeMapping.java new file mode 100644 index 0000000000..76c7fc6eed --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ProxyNodeMapping.java @@ -0,0 +1,404 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.server.MetadataCache; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.message.MetadataResponseData; +import org.apache.kafka.image.MetadataDelta; +import org.apache.kafka.image.MetadataImage; +import org.apache.kafka.metadata.BrokerRegistration; + +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import software.amazon.awssdk.annotations.NotNull; +import thirdparty.com.github.jaskey.consistenthash.ConsistentHashRouter; + +/** + * Maintain the relationship for main node and proxy node. + */ +class ProxyNodeMapping { + private static final Logger LOGGER = LoggerFactory.getLogger(ProxyNodeMapping.class); + private static final String NOOP_RACK = ""; + private static final int DEFAULT_VIRTUAL_NODE_COUNT = 8; + private final Node currentNode; + private final String currentRack; + private final String interBrokerListenerName; + private final MetadataCache metadataCache; + private final List listeners = new CopyOnWriteArrayList<>(); + + volatile Map> main2proxyByRack = new HashMap<>(); + volatile boolean inited = false; + + public ProxyNodeMapping(Node currentNode, String currentRack, String interBrokerListenerName, + MetadataCache metadataCache) { + this.interBrokerListenerName = interBrokerListenerName; + this.currentNode = currentNode; + this.currentRack = currentRack; + this.metadataCache = metadataCache; + Threads.COMMON_SCHEDULER.scheduleWithFixedDelay(() -> logMapping(main2proxyByRack), 1, 1, TimeUnit.MINUTES); + } + + /** + * Get route out node to split the produce request. + *

+ * If return Node.noNode, it means the producer should refresh metadata and send to another node, {@link RouterOut} will return NOT_LEADER_OR_FOLLOWER. + */ + public Node getRouteOutNode(String topicName, int partition, ClientIdMetadata clientId) { + String clientRack = clientId.rack(); + + BrokerRegistration target = metadataCache.getPartitionLeaderNode(topicName, partition); + if (target == null) { + return currentNode; + } + if (clientRack == null) { + // If the client rack isn't set, expect produce send to the real leader. + if (target.id() == currentNode.id()) { + return currentNode; + } else { + return Node.noNode(); + } + } + + Map> main2proxyByRack = this.main2proxyByRack; + if (Objects.equals(clientRack, currentRack)) { + if (target.id() == currentNode.id()) { + return currentNode; + } else { + if (Objects.equals(target.rack().orElse(null), currentRack)) { + // The producer should refresh metadata and send to another node in the same rack as the producer + return Node.noNode(); + } else { + // Check whether the current node should proxy the target + Map main2proxy = main2proxyByRack.get(currentRack); + if (main2proxy == null) { + // The current node is the last node in the rack, and the current node is shutting down. + return Node.noNode(); + } + BrokerRegistration proxyNode = main2proxy.get(target.id()); + if (proxyNode != null && proxyNode.id() == currentNode.id()) { + // Get the target main node. + return target.node(interBrokerListenerName).orElse(currentNode); + } else { + // The producer should refresh metadata and send to another node in the same rack as the current node. + return Node.noNode(); + } + + } + } + } else { + if (main2proxyByRack.containsKey(clientRack)) { + // The producer should send records to the nodes with the same rack. + return Node.noNode(); + } else { + MismatchRecorder.instance().record(topicName, clientId); + // The cluster doesn't cover the client rack, the producer should directly send records to the partition main node. + if (target.id() == currentNode.id()) { + return currentNode; + } else { + return Node.noNode(); + } + } + } + } + + /** + * Get the proxy leader node when NOT_LEADER_OR_FOLLOWER happens. + */ + public Optional getLeaderNode(int leaderMainNodeId, ClientIdMetadata clientId, String listenerName) { + BrokerRegistration target = metadataCache.getNode(leaderMainNodeId); + if (target == null) { + return Optional.empty(); + } + String clientRack = clientId.rack(); + if (clientRack == null) { + // If the client rack isn't set, then return the main node. + return target.node(listenerName); + } + + Map clientRackMain2proxy = main2proxyByRack.get(clientRack); + if (clientRackMain2proxy == null) { + // If the cluster doesn't cover the client rack, the producer should directly send records to the main node. + return target.node(listenerName); + } + + // Get the proxy node. + BrokerRegistration proxy = clientRackMain2proxy.get(target.id()); + if (proxy == null) { + // The producer rack is the same as the leader rack. + return target.node(listenerName); + } + return proxy.node(listenerName); + } + + public List handleMetadataResponse(ClientIdMetadata clientIdMetadata, + List topics) { + String clientRack = clientIdMetadata.rack(); + if (clientRack == null) { + return withSnapshotReadFollowers(topics); + } + Map clientRackMain2proxy = main2proxyByRack.get(clientRack); + if (clientRackMain2proxy == null) { + // If the cluster doesn't cover the client rack, the producer should directly send records to the main node. + return withSnapshotReadFollowers(topics); + } + // If the client config rack in clientId, we need to replace the leader id with the proxy leader id. + topics.forEach(metadataResponseTopic -> { + metadataResponseTopic.partitions().forEach(metadataResponsePartition -> { + int leaderMainNodeId = metadataResponsePartition.leaderId(); + if (leaderMainNodeId != -1) { + BrokerRegistration proxy = clientRackMain2proxy.get(leaderMainNodeId); + if (proxy != null) { + int proxyLeaderId = proxy.id(); + if (proxyLeaderId != leaderMainNodeId) { + metadataResponsePartition.setLeaderId(proxyLeaderId); + metadataResponsePartition.setIsrNodes(List.of(proxyLeaderId)); + metadataResponsePartition.setReplicaNodes(List.of(proxyLeaderId)); + } + } + } + }); + }); + return topics; + } + + public void onChange(MetadataDelta delta, MetadataImage image) { + if (!inited) { + // When the main2proxyByRack is un-inited, we should force update. + inited = true; + } else { + if (delta.clusterDelta() == null || delta.clusterDelta().changedBrokers().isEmpty()) { + return; + } + } + // categorize the brokers by rack + Map> rack2brokers = new HashMap<>(); + image.cluster().brokers().forEach((nodeId, node) -> { + if (node.fenced() || node.inControlledShutdown()) { + return; + } + rack2brokers.compute(node.rack().orElse(NOOP_RACK), (rack, list) -> { + if (list == null) { + list = new ArrayList<>(); + } + list.add(node); + return list; + }); + }); + this.main2proxyByRack = calMain2proxyByRack(rack2brokers); + logMapping(main2proxyByRack); + notifyListeners(this.main2proxyByRack); + } + + public void registerListener(ProxyTopologyChangeListener listener) { + listeners.add(listener); + listener.onChange(this.main2proxyByRack); + } + + private void notifyListeners(Map> main2proxyByRack) { + listeners.forEach(listener -> { + try { + listener.onChange(main2proxyByRack); + } catch (Throwable e) { + LOGGER.error("fail to notify listener {}", listener, e); + } + }); + } + + private List withSnapshotReadFollowers(List topics) { + topics.forEach(metadataResponseTopic -> { + metadataResponseTopic.partitions().forEach(metadataResponsePartition -> { + int leaderMainNodeId = metadataResponsePartition.leaderId(); + if (leaderMainNodeId != -1) { + List replicas = new ArrayList<>(main2proxyByRack.size()); + replicas.add(leaderMainNodeId); + main2proxyByRack.forEach((rack, main2proxy) -> { + BrokerRegistration proxy = main2proxy.get(leaderMainNodeId); + if (proxy != null && proxy.id() != leaderMainNodeId) { + replicas.add(proxy.id()); + } + }); + metadataResponsePartition.setIsrNodes(replicas); + metadataResponsePartition.setReplicaNodes(replicas); + } + }); + }); + return topics; + } + + static Map> calMain2proxyByRack( + Map> rack2brokers) { + rack2brokers.forEach((rack, brokers) -> brokers.sort(Comparator.comparingInt(BrokerRegistration::id))); + + Map> newMain2proxyByRack = new HashMap<>(); + List racks = rack2brokers.keySet().stream().sorted().collect(Collectors.toList()); + racks.forEach(proxyRack -> { + Map newMain2proxy = new HashMap<>(); + List proxyNodes = new ArrayList<>(); + ConsistentHashRouter router = new ConsistentHashRouter<>(); + List proxyRackBrokers = rack2brokers.get(proxyRack); + proxyRackBrokers.forEach(node -> { + ProxyNode proxyNode = new ProxyNode(node); + router.addNode(proxyNode, DEFAULT_VIRTUAL_NODE_COUNT); + proxyNodes.add(proxyNode); + }); + + // allocate the proxy node by consistent hash + int proxyNodeCount = 0; + for (String rack : racks) { + List brokers = rack2brokers.get(rack); + if (Objects.equals(rack, proxyRack)) { + continue; + } + for (BrokerRegistration node : brokers) { + ProxyNode proxyNode = router.routeNode(Integer.toString(node.id())); + newMain2proxy.put(node.id(), proxyNode.node); + proxyNode.mainNodeIds.add(node.id()); + proxyNodeCount++; + } + } + // balance the proxy node count + double avg = Math.ceil((double) proxyNodeCount / proxyNodes.size()); + proxyNodes.sort(Comparator.reverseOrder()); + for (ProxyNode overloadNode : proxyNodes) { + if (overloadNode.mainNodeIds.size() <= avg) { + break; + } + // move overload node's proxied node to free node + for (int i = proxyNodes.size() - 1; i >= 0 && overloadNode.mainNodeIds.size() > avg; i--) { + ProxyNode freeNode = proxyNodes.get(i); + if (freeNode.mainNodeIds.size() > avg - 1) { + continue; + } + Integer mainNodeId = overloadNode.mainNodeIds.remove(overloadNode.mainNodeIds.size() - 1); + newMain2proxy.put(mainNodeId, freeNode.node); + freeNode.mainNodeIds.add(mainNodeId); + } + } + // try let controller only proxy controller + tryFreeController(proxyNodes, avg); + + newMain2proxyByRack.put(proxyRack, newMain2proxy); + }); + return newMain2proxyByRack; + } + + /** + * Try to move the traffic from controller to broker. + * - Let main node(controller) proxied by proxy node(controller). + * - Let proxy node(controller) proxy less main node if possible. + */ + static void tryFreeController(List proxyNodes, double avg) { + for (ProxyNode controller : proxyNodes) { + if (!isController(controller.node.id())) { + continue; + } + for (int i = 0; i < controller.mainNodeIds.size(); i++) { + int mainNodeId = controller.mainNodeIds.get(i); + if (isController(mainNodeId)) { + continue; + } + L1: + for (ProxyNode switchNode : proxyNodes) { + if (switchNode.node.id() == controller.node.id() || isController(switchNode.node.id())) { + continue; + } + // move the main node to the switch node + if (switchNode.mainNodeIds.size() < avg) { + controller.mainNodeIds.remove(i); + switchNode.mainNodeIds.add(mainNodeId); + i--; + break; + } else { + // swap the main node with the switch node's main node(controller) + for (int j = 0; j < switchNode.mainNodeIds.size(); j++) { + int switchNodeMainNodeId = switchNode.mainNodeIds.get(j); + if (!isController(switchNodeMainNodeId)) { + continue; + } + controller.mainNodeIds.set(i, switchNodeMainNodeId); + switchNode.mainNodeIds.set(j, mainNodeId); + break L1; + } + } + } + } + } + } + + + + static void logMapping(Map> main2proxyByRack) { + StringBuilder sb = new StringBuilder(); + main2proxyByRack.forEach((rack, main2proxy) -> + main2proxy.forEach((mainNodeId, proxyNode) -> + sb.append(" Main ").append(mainNodeId).append(" => Proxy ").append(proxyNode.id()).append("(").append(rack).append(")\n") + ) + ); + LOGGER.info("ProxyNodeMapping:\n{}", sb); + } + + static class ProxyNode implements thirdparty.com.github.jaskey.consistenthash.Node, Comparable { + final BrokerRegistration node; + final List mainNodeIds = new ArrayList<>(); + + private final String key; + + public ProxyNode(BrokerRegistration node) { + this.node = node; + this.key = Integer.toString(node.id()); + } + + @Override + public String getKey() { + return key; + } + + @Override + public int compareTo(@NotNull ProxyNode o) { + int rst = Integer.compare(mainNodeIds.size(), o.mainNodeIds.size()); + if (rst != 0) { + return rst; + } + return Integer.compare(node.id(), o.node.id()); + } + } + + static boolean isController(int nodeId) { + return nodeId < 100; + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/ProxyTopologyChangeListener.java b/core/src/main/java/kafka/automq/zerozone/ProxyTopologyChangeListener.java new file mode 100644 index 0000000000..1de256432b --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ProxyTopologyChangeListener.java @@ -0,0 +1,30 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.metadata.BrokerRegistration; + +import java.util.Map; + +public interface ProxyTopologyChangeListener { + + void onChange(Map> main2proxyByRack); + +} diff --git a/core/src/main/java/kafka/automq/zerozone/Replayer.java b/core/src/main/java/kafka/automq/zerozone/Replayer.java new file mode 100644 index 0000000000..1dd815e516 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/Replayer.java @@ -0,0 +1,43 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.WriteAheadLog; + +import java.util.List; +import java.util.concurrent.CompletableFuture; + +public interface Replayer { + + /** + * Replay SSO to snapshot-read cache + */ + CompletableFuture replay(List objects); + + /** + * Replay WAL to snapshot-read cache. + * If the record in WAL is a linked record, it will decode the linked record to the real record. + */ + CompletableFuture replay(WriteAheadLog confirmWAL, RecordOffset startOffset, RecordOffset endOffset, List walRecords); + +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterChannel.java b/core/src/main/java/kafka/automq/zerozone/RouterChannel.java new file mode 100644 index 0000000000..99ab5a99ad --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterChannel.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import java.util.concurrent.CompletableFuture; + +import io.netty.buffer.ByteBuf; + +public interface RouterChannel { + + CompletableFuture append(int targetNodeId, short orderHint, ByteBuf data); + + CompletableFuture get(ByteBuf channelOffset); + + void nextEpoch(long epoch); + + void trim(long epoch); + + CompletableFuture close(); + + class AppendResult { + private final long epoch; + private final ByteBuf channelOffset; + + public AppendResult(long epoch, ByteBuf channelOffset) { + this.epoch = epoch; + this.channelOffset = channelOffset; + } + + public long epoch() { + return epoch; + } + + public ByteBuf channelOffset() { + return channelOffset; + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterChannelProvider.java b/core/src/main/java/kafka/automq/zerozone/RouterChannelProvider.java new file mode 100644 index 0000000000..1013b29328 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterChannelProvider.java @@ -0,0 +1,40 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.controller.stream.RouterChannelEpoch; +import org.apache.kafka.image.loader.MetadataListener; + +public interface RouterChannelProvider extends MetadataListener { + + RouterChannel channel(); + + RouterChannel readOnlyChannel(int node); + + RouterChannelEpoch epoch(); + + void addEpochListener(EpochListener listener); + + void close(); + + interface EpochListener { + void onNewEpoch(RouterChannelEpoch epoch); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterIn.java b/core/src/main/java/kafka/automq/zerozone/RouterIn.java new file mode 100644 index 0000000000..0e5ad58b65 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterIn.java @@ -0,0 +1,184 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdKey; +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.automq.interceptor.ProduceRequestArgs; +import kafka.server.RequestLocal; +import kafka.server.streamaspect.ElasticKafkaApis; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.errors.ThrottlingQuotaExceededException; +import org.apache.kafka.common.message.AutomqZoneRouterResponseData; +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.MutableRecordBatch; +import org.apache.kafka.common.requests.ProduceResponse; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; + +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Semaphore; + +import io.netty.buffer.Unpooled; +import io.netty.util.concurrent.FastThreadLocal; + +class RouterIn { + private static final Logger LOGGER = LoggerFactory.getLogger(RouterIn.class); + + private final Semaphore inflightAppendLimiter = new Semaphore(Math.min(50 * 1024 * 1024 * Systems.CPU_CORES, Integer.MAX_VALUE)); + private CompletableFuture lastRouterCf = CompletableFuture.completedFuture(null); + + private final ExecutorService executor = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(1, "object-cross-zone-produce-router-in", true, LOGGER); + private final FastThreadLocal requestLocals = new FastThreadLocal<>() { + @Override + protected RequestLocal initialValue() { + // The same as KafkaRequestHandler.requestLocal + return RequestLocal.withThreadConfinedCaching(); + } + }; + + private RouterInProduceHandler routerInProduceHandler; + + private final ObjectStorage objectStorage; + private final ElasticKafkaApis kafkaApis; + private final String rack; + + public RouterIn(ObjectStorage objectStorage, ElasticKafkaApis kafkaApis, String rack) { + this.objectStorage = objectStorage; + this.kafkaApis = kafkaApis; + this.rack = rack; + this.routerInProduceHandler = kafkaApis::handleProduceAppendJavaCompatible; + } + + public synchronized CompletableFuture handleZoneRouterRequest(byte[] metadata) { + RouterRecord routerRecord = RouterRecord.decode(Unpooled.wrappedBuffer(metadata)); + try { + if (inflightAppendLimiter.tryAcquire(routerRecord.size())) { + return handleZoneRouterRequest0(routerRecord).whenComplete((rst, ex) -> { + inflightAppendLimiter.release(routerRecord.size()); + }); + } else { + return FutureUtil.failedFuture(new ThrottlingQuotaExceededException("inflight append limit exceeded")); + } + } catch (Throwable e) { + inflightAppendLimiter.release(routerRecord.size()); + LOGGER.error("[UNEXPECTED] handleZoneRouterRequest failed", e); + return FutureUtil.failedFuture(e); + } + } + + public CompletableFuture handleZoneRouterRequest0(RouterRecord routerRecord) { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ROUTER_IN],data={}", routerRecord); + } + // inbound is consumed by the object storage + ZeroZoneMetricsManager.recordRouterInBytes(routerRecord.nodeId(), routerRecord.size()); + CompletableFuture> readCf = new ZoneRouterPackReader(routerRecord.nodeId(), routerRecord.bucketId(), routerRecord.objectId(), objectStorage) + .readProduceRequests(new Position(routerRecord.position(), routerRecord.size())); + // Orderly handle the request + CompletableFuture prevLastRouterCf = lastRouterCf; + CompletableFuture appendCf = readCf + .thenCompose(rst -> prevLastRouterCf.thenApply(nil -> rst)) + .thenComposeAsync(produces -> { + List> cfList = new ArrayList<>(); + produces.stream().map(this::append).forEach(cfList::add); + return CompletableFuture.allOf(cfList.toArray(new CompletableFuture[0])).thenApply(nil -> { + AutomqZoneRouterResponseData response = new AutomqZoneRouterResponseData(); + cfList.forEach(cf -> response.responses().add(cf.join())); + return new AutomqZoneRouterResponse(response); + }); + }, executor); + this.lastRouterCf = appendCf.thenAccept(rst -> { + }).exceptionally(ex -> null); + return appendCf; + } + + public void setRouterInProduceHandler(RouterInProduceHandler routerInProduceHandler) { + this.routerInProduceHandler = routerInProduceHandler; + } + + CompletableFuture append( + ZoneRouterProduceRequest zoneRouterProduceRequest) { + ZoneRouterProduceRequest.Flag flag = new ZoneRouterProduceRequest.Flag(zoneRouterProduceRequest.flag()); + ProduceRequestData data = zoneRouterProduceRequest.data(); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("[ROUTER_IN],[APPEND],data={}", data); + } + + Map realEntriesPerPartition = ZeroZoneTrafficInterceptor.produceRequestToMap(data); + short apiVersion = zoneRouterProduceRequest.apiVersion(); + CompletableFuture cf = new CompletableFuture<>(); + cf.whenComplete((resp, ex) -> zoneRouterProduceRequest.close()); + // TODO: parallel request for different partitions + routerInProduceHandler.handleProduceAppend( + ProduceRequestArgs.builder() + .clientId(buildClientId(realEntriesPerPartition)) + .timeout(10000) + .requiredAcks(data.acks()) + .internalTopicsAllowed(flag.internalTopicsAllowed()) + .transactionId(data.transactionalId()) + .entriesPerPartition(realEntriesPerPartition) + .responseCallback(rst -> { + @SuppressWarnings("deprecation") + ProduceResponse produceResponse = new ProduceResponse(rst, 0, Collections.emptyList()); + AutomqZoneRouterResponseData.Response response = new AutomqZoneRouterResponseData.Response() + .setData(ZoneRouterResponseCodec.encode(produceResponse.data()).array()); + cf.complete(response); + }) + .recordValidationStatsCallback(rst -> { + }) + .apiVersion(apiVersion) + .requestLocal(requestLocals.get()) + .build() + ); + return cf; + } + + protected ClientIdMetadata buildClientId(Map entriesPerPartition) { + String clientId = String.format("%s=%s", ClientIdKey.AVAILABILITY_ZONE, rack); + String connectionId = getConnectionIdFrom(entriesPerPartition); + return ClientIdMetadata.of(clientId, null, connectionId); + } + + protected String getConnectionIdFrom(Map entriesPerPartition) { + for (Map.Entry entry : entriesPerPartition.entrySet()) { + for (MutableRecordBatch batch : entry.getValue().batches()) { + if (batch.hasProducerId()) { + return String.valueOf(batch.producerId()); + } + } + } + return null; + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterInProduceHandler.java b/core/src/main/java/kafka/automq/zerozone/RouterInProduceHandler.java new file mode 100644 index 0000000000..8f62cfa5b5 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterInProduceHandler.java @@ -0,0 +1,26 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ProduceRequestArgs; + +public interface RouterInProduceHandler { + void handleProduceAppend(ProduceRequestArgs args); +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/zerozone/RouterInV2.java b/core/src/main/java/kafka/automq/zerozone/RouterInV2.java new file mode 100644 index 0000000000..f4827b936d --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterInV2.java @@ -0,0 +1,266 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdKey; +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.automq.interceptor.ProduceRequestArgs; +import kafka.server.KafkaRequestHandler; +import kafka.server.RequestLocal; +import kafka.server.streamaspect.ElasticKafkaApis; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.message.AutomqZoneRouterRequestData; +import org.apache.kafka.common.message.AutomqZoneRouterResponseData; +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.MutableRecordBatch; +import org.apache.kafka.common.requests.ProduceResponse; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; +import org.apache.kafka.common.utils.Time; + +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.threads.EventLoop; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.util.concurrent.FastThreadLocal; + +public class RouterInV2 implements NonBlockingLocalRouterHandler { + private static final Logger LOGGER = LoggerFactory.getLogger(RouterInV2.class); + + static { + // RouterIn will parallel append the records from one AutomqZoneRouterRequest. + // So the append thread isn't the KafkaRequestHandler + KafkaRequestHandler.setBypassThreadCheck(true); + } + + private final RouterChannelProvider channelProvider; + private final ElasticKafkaApis kafkaApis; + private final String rack; + private final RouterInProduceHandler localAppendHandler; + private RouterInProduceHandler routerInProduceHandler; + private final BlockingQueue unpackLinkQueue = new ArrayBlockingQueue<>(Systems.CPU_CORES * 8192); + private final EventLoop[] appendEventLoops; + private final FastThreadLocal requestLocals = new FastThreadLocal<>() { + @Override + protected RequestLocal initialValue() { + // The same as KafkaRequestHandler.requestLocal + return RequestLocal.withThreadConfinedCaching(); + } + }; + private final Time time; + + public RouterInV2(RouterChannelProvider channelProvider, ElasticKafkaApis kafkaApis, String rack, Time time) { + this.channelProvider = channelProvider; + this.kafkaApis = kafkaApis; + this.rack = rack; + this.localAppendHandler = kafkaApis::handleProduceAppendJavaCompatible; + this.routerInProduceHandler = this.localAppendHandler; + this.time = time; + + this.appendEventLoops = new EventLoop[Systems.CPU_CORES]; + for (int i = 0; i < appendEventLoops.length; i++) { + this.appendEventLoops[i] = new EventLoop("ROUTER_IN_V2_APPEND_" + i); + } + + } + + public CompletableFuture handleZoneRouterRequest(AutomqZoneRouterRequestData request) { + long requestEpoch = request.routeEpoch(); + if (requestEpoch <= channelProvider.epoch().getFenced()) { + String str = String.format("The router request epoch %s is less than fenced epoch %s.", requestEpoch, channelProvider.epoch().getFenced()); + return CompletableFuture.failedFuture(new IllegalStateException(str)); + } + return handleZoneRouterRequest0(request); + } + + private CompletableFuture handleZoneRouterRequest0(AutomqZoneRouterRequestData request) { + RouterRecordV2 routerRecord = RouterRecordV2.decode(Unpooled.wrappedBuffer(request.metadata())); + RouterChannel routerChannel = channelProvider.readOnlyChannel(routerRecord.nodeId()); + List> subResponseList = new ArrayList<>(routerRecord.channelOffsets().size()); + AtomicInteger size = new AtomicInteger(0); + long startNanos = time.nanoseconds(); + for (ByteBuf channelOffset : routerRecord.channelOffsets()) { + PartitionProduceRequest partitionProduceRequest = new PartitionProduceRequest(ChannelOffset.of(channelOffset)); + partitionProduceRequest.unpackLinkCf = routerChannel.get(channelOffset); + addToUnpackLinkQueue(partitionProduceRequest); + partitionProduceRequest.unpackLinkCf.whenComplete((rst, ex) -> { + if (ex == null) { + size.addAndGet(rst.readableBytes()); + } + handleUnpackLink(); + ZeroZoneMetricsManager.GET_CHANNEL_LATENCY.record(time.nanoseconds() - startNanos); + }); + subResponseList.add(partitionProduceRequest.responseCf); + } + return CompletableFuture.allOf(subResponseList.toArray(new CompletableFuture[0])).thenApply(nil -> { + AutomqZoneRouterResponseData response = new AutomqZoneRouterResponseData(); + response.setResponses(subResponseList.stream().map(CompletableFuture::join).collect(Collectors.toList())); + ZeroZoneMetricsManager.recordRouterInBytes(routerRecord.nodeId(), size.get()); + return new AutomqZoneRouterResponse(response); + }); + } + + private void handleUnpackLink() { + if (unpackLinkQueue.isEmpty()) { + return; + } + synchronized (unpackLinkQueue) { + while (!unpackLinkQueue.isEmpty()) { + PartitionProduceRequest req = unpackLinkQueue.peek(); + if (req.unpackLinkCf.isDone()) { + EventLoop eventLoop = appendEventLoops[Math.abs(req.channelOffset.orderHint() % appendEventLoops.length)]; + req.unpackLinkCf.thenComposeAsync(buf -> { + ZoneRouterProduceRequest zoneRouterProduceRequest = ZoneRouterPackReader.decodeDataBlock(buf).get(0); + try { + return append0(req.channelOffset, zoneRouterProduceRequest, false); + } finally { + buf.release(); + } + }, eventLoop).whenComplete((resp, ex) -> { + if (ex != null) { + LOGGER.error("[ROUTER_IN],[FAILED]", ex); + req.responseCf.completeExceptionally(ex); + return; + } + req.responseCf.complete(resp); + }); + unpackLinkQueue.poll(); + } else { + break; + } + } + } + } + + private void addToUnpackLinkQueue(PartitionProduceRequest req) { + for (;;) { + try { + unpackLinkQueue.put(req); + return; + } catch (InterruptedException ignored) { + } + } + } + + @Override + public CompletableFuture append( + ChannelOffset channelOffset, + ZoneRouterProduceRequest zoneRouterProduceRequest + ) { + CompletableFuture cf = new CompletableFuture<>(); + appendEventLoops[Math.abs(channelOffset.orderHint() % appendEventLoops.length)].execute(() -> + FutureUtil.propagate(append0(channelOffset, zoneRouterProduceRequest, true), cf)); + return cf; + } + + private CompletableFuture append0( + ChannelOffset channelOffset, + ZoneRouterProduceRequest zoneRouterProduceRequest, + boolean local + ) { + ZoneRouterProduceRequest.Flag flag = new ZoneRouterProduceRequest.Flag(zoneRouterProduceRequest.flag()); + ProduceRequestData data = zoneRouterProduceRequest.data(); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ROUTER_IN],[APPEND],data={}", data); + } + + ZeroZoneThreadLocalContext.writeContext().channelOffset = channelOffset; + Map realEntriesPerPartition = ZeroZoneTrafficInterceptor.produceRequestToMap(data); + short apiVersion = zoneRouterProduceRequest.apiVersion(); + CompletableFuture cf = new CompletableFuture<>(); + RouterInProduceHandler handler = local ? localAppendHandler : routerInProduceHandler; + // We should release the request after append completed. + cf.whenComplete((resp, ex) -> FutureUtil.suppress(zoneRouterProduceRequest::close, LOGGER)); + handler.handleProduceAppend( + ProduceRequestArgs.builder() + .clientId(buildClientId(realEntriesPerPartition)) + .timeout(data.timeoutMs()) + // The CommittedEpochManager requires the data to be persisted prior to bumping the committed epoch. + .requiredAcks((short) -1) + .internalTopicsAllowed(flag.internalTopicsAllowed()) + .transactionId(data.transactionalId()) + .entriesPerPartition(realEntriesPerPartition) + .responseCallback(rst -> { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ROUTER_IN],[RESPONSE],response={}", rst); + } + @SuppressWarnings("deprecation") + ProduceResponse produceResponse = new ProduceResponse(rst, 0, Collections.emptyList()); + AutomqZoneRouterResponseData.Response response = new AutomqZoneRouterResponseData.Response() + .setData(ZoneRouterResponseCodec.encode(produceResponse.data()).array()); + cf.complete(response); + }) + .recordValidationStatsCallback(rst -> { + }) + .apiVersion(apiVersion) + .requestLocal(requestLocals.get()) + .build() + ); + return cf; + } + + public void setRouterInProduceHandler(RouterInProduceHandler routerInProduceHandler) { + this.routerInProduceHandler = routerInProduceHandler; + } + + protected ClientIdMetadata buildClientId(Map entriesPerPartition) { + String clientId = String.format("%s=%s", ClientIdKey.AVAILABILITY_ZONE, rack); + String connectionId = getConnectionIdFrom(entriesPerPartition); + return ClientIdMetadata.of(clientId, null, connectionId); + } + + protected String getConnectionIdFrom(Map entriesPerPartition) { + for (Map.Entry entry : entriesPerPartition.entrySet()) { + for (MutableRecordBatch batch : entry.getValue().batches()) { + if (batch.hasProducerId()) { + return String.valueOf(batch.producerId()); + } + } + } + return null; + } + + static class PartitionProduceRequest { + final ChannelOffset channelOffset; + CompletableFuture unpackLinkCf; + final CompletableFuture responseCf = new CompletableFuture<>(); + + public PartitionProduceRequest(ChannelOffset channelOffset) { + this.channelOffset = channelOffset; + } + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterOut.java b/core/src/main/java/kafka/automq/zerozone/RouterOut.java new file mode 100644 index 0000000000..d4ab2533df --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterOut.java @@ -0,0 +1,434 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.automq.interceptor.ProduceRequestArgs; +import kafka.server.streamaspect.ElasticKafkaApis; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.message.AutomqZoneRouterRequestData; +import org.apache.kafka.common.message.AutomqZoneRouterResponseData; +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.message.ProduceResponseData; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.requests.ProduceResponse; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterRequest; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; +import org.apache.kafka.common.utils.Time; + +import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; +import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; +import com.automq.stream.s3.network.NetworkBandwidthLimiter; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.ThreadUtils; +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; + +import io.netty.buffer.Unpooled; +import io.netty.util.HashedWheelTimer; +import io.netty.util.Timeout; + +public class RouterOut { + private static final Logger LOGGER = LoggerFactory.getLogger(RouterOut.class); + private static final String BATCH_INTERVAL_CONFIG = "batchInterval"; + private static final String MAX_BYTES_IN_BATCH_CONFIG = "maxBytesInBatch"; + + private final int batchIntervalMs; + private final int batchSizeThreshold; + private final Node currentNode; + + private final Map> pendingRequests = new ConcurrentHashMap<>(); + private CompletableFuture lastRouterCf = CompletableFuture.completedFuture(null); + private final AtomicInteger batchSize = new AtomicInteger(); + private final AtomicLong nextObjectId = new AtomicLong(); + private long lastUploadTimestamp = 0; + + private final ScheduledExecutorService scheduler = Threads.newSingleThreadScheduledExecutor("object-cross-zone-produce-router-out", true, LOGGER); + private final HashedWheelTimer timeoutDetect = new HashedWheelTimer( + ThreadUtils.createThreadFactory("object-cross-zone-produce-router-out-timeout-detect", true), + 1, TimeUnit.SECONDS, 100 + ); + + private final AsyncSender asyncSender; + private final BucketURI bucketURI; + private final ObjectStorage objectStorage; + private final GetRouterOutNode mapping; + private final ElasticKafkaApis kafkaApis; + private final Time time; + private final NetworkBandwidthLimiter inboundLimiter = GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.INBOUND); + + private final PerfMode perfMode = new PerfMode(); + + public RouterOut(Node currentNode, BucketURI bucketURI, ObjectStorage objectStorage, + GetRouterOutNode mapping, ElasticKafkaApis kafkaApis, AsyncSender asyncSender, Time time) { + this.currentNode = currentNode; + this.bucketURI = bucketURI; + this.batchIntervalMs = Integer.parseInt(bucketURI.extensionString(BATCH_INTERVAL_CONFIG, "250")); + this.batchSizeThreshold = Integer.parseInt(bucketURI.extensionString(MAX_BYTES_IN_BATCH_CONFIG, "8388608")); + this.objectStorage = objectStorage; + this.mapping = mapping; + this.kafkaApis = kafkaApis; + this.asyncSender = asyncSender; + this.time = time; + cleanup(); + scheduler.scheduleWithFixedDelay(this::proxy, batchIntervalMs, batchIntervalMs, TimeUnit.MILLISECONDS); + } + + public void handleProduceAppendProxy(ProduceRequestArgs args) { + short flag = new ZoneRouterProduceRequest.Flag().internalTopicsAllowed(args.internalTopicsAllowed()).value(); + Map requests = split(args.apiVersion(), args.clientId(), args.timeout(), flag, args.requiredAcks(), args.transactionId(), args.entriesPerPartition()); + + boolean forceRoute = perfMode.isRouteInPerfMode(args.entriesPerPartition()); + requests.forEach((node, request) -> { + if (node.id() == Node.noNode().id()) { + request.completeWithNotLeaderNotFollower(); + return; + } + if (node.id() == currentNode.id() && !forceRoute) { + kafkaApis.handleProduceAppendJavaCompatible( + args.toBuilder() + .entriesPerPartition(ZeroZoneTrafficInterceptor.produceRequestToMap(request.data)) + .responseCallback(responseCallbackRst -> { + request.cf.complete(responseCallbackRst); + }) + .build() + ); + } else { + ZeroZoneMetricsManager.recordRouterOutBytes(node.id(), request.size); + pendingRequests.compute(node, (n, queue) -> { + if (queue == null) { + queue = new LinkedBlockingQueue<>(); + } + queue.add(request); + batchSize.addAndGet(request.size); + return queue; + }); + } + }); + + Timeout timeout = timeoutDetect.newTimeout(t -> LOGGER.error("[POTENTIAL_BUG] router out timeout, {}", t), 1, TimeUnit.MINUTES); + // Group the request result + Map rst = new ConcurrentHashMap<>(); + List> proxyRequestCfList = requests.values().stream().map( + request -> request.cf + .thenAccept(rst::putAll) + .exceptionally(ex -> { + LOGGER.error("[UNEXPECTED],request={}", request.topicPartitions, ex); + request.topicPartitions.forEach(topicPartition -> + rst.put(topicPartition, new ProduceResponse.PartitionResponse(Errors.UNKNOWN_SERVER_ERROR))); + return null; + }) + ).collect(Collectors.toList()); + CompletableFuture.allOf(proxyRequestCfList.toArray(CompletableFuture[]::new)) + .thenAccept(nil -> args.responseCallback().accept(rst)) + .whenComplete((nil, ex) -> { + if (timeout.isExpired()) { + LOGGER.error("[POTENTIAL_BUG_RECOVERED] router out timeout recover, {}", timeout); + } else { + timeout.cancel(); + } + }); + + if (batchSize.get() >= batchSizeThreshold || time.milliseconds() - batchIntervalMs >= lastUploadTimestamp) { + scheduler.submit(this::proxy); + } + } + + private void proxy() { + try { + proxy0(); + } catch (Throwable e) { + LOGGER.error("[UNEXPECTED],[BUG],proxy failed", e); + } + } + + private void proxy0() { + if (batchSize.get() < batchSizeThreshold && time.milliseconds() - batchIntervalMs < lastUploadTimestamp) { + return; + } + + lastUploadTimestamp = time.milliseconds(); + + // 1. Batch the request by destination node. + Map> node2requests = new HashMap<>(); + pendingRequests.forEach((node, queue) -> { + List requests = new ArrayList<>(); + queue.drainTo(requests); + if (!requests.isEmpty()) { + node2requests.put(node, requests); + batchSize.addAndGet(-requests.stream().mapToInt(r -> r.size).sum()); + } + }); + if (node2requests.isEmpty()) { + return; + } + + // 2. Write the request to object and get the relative position. + long objectId = nextObjectId.incrementAndGet(); + ZoneRouterPackWriter writer = new ZoneRouterPackWriter(currentNode.id(), objectId, objectStorage); + Map node2position = new HashMap<>(); + node2requests.forEach((node, requests) -> { + Position position = writer + .addProduceRequests( + requests + .stream() + .map(r -> new ZoneRouterProduceRequest(r.apiVersion, r.flag, r.data)) + .collect(Collectors.toList()) + ); + requests.forEach(ProxyRequest::afterRouter); + node2position.put(node, position); + }); + + // 3. send ZoneRouterRequest to the route out nodes. + CompletableFuture writeCf = writer.close(); + CompletableFuture prevLastRouterCf = lastRouterCf; + lastRouterCf = writeCf + // Orderly send the router request. + .thenCompose(nil -> prevLastRouterCf) + .thenAccept(nil -> { + // TODO: in / out metrics and type (要经过限流器 + 另外一个 type 类型统计) + List> sendCfList = new ArrayList<>(); + node2position.forEach((destNode, position) -> { + List proxyRequests = node2requests.get(destNode); + CompletableFuture sendCf = sendRouterRequest(destNode, objectId, position, proxyRequests); + sendCfList.add(sendCf); + }); + CompletableFuture.allOf(sendCfList.toArray(new CompletableFuture[0])).whenComplete((nil2, ex) -> { + ObjectStorage.ObjectPath path = writer.objectPath(); + objectStorage.delete(List.of(path)).exceptionally(ex2 -> { + LOGGER.error("delete {} fail", path, ex); + return null; + }); + }); + }) + .exceptionally(ex -> { + LOGGER.error("[UNEXPECTED],[PROXY]", ex); + return null; + }); + } + + private CompletableFuture sendRouterRequest(Node destNode, long objectId, Position position, + List requests) { + RouterRecord routerRecord = new RouterRecord(currentNode.id(), objectStorage.bucketId(), objectId, position.position(), position.size()); + + AutomqZoneRouterRequest.Builder builder = new AutomqZoneRouterRequest.Builder( + new AutomqZoneRouterRequestData().setMetadata(routerRecord.encode().array()) + ); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ROUTER_OUT],node={},objectId={},position={},partitions={}", destNode, objectId, position, requests.stream().map(r -> r.topicPartitions).collect(Collectors.toList())); + } + return asyncSender.sendRequest(destNode, builder).thenAccept(clientResponse -> { + if (!clientResponse.hasResponse()) { + LOGGER.error("[ROUTER_OUT],[NO_RESPONSE],response={}", clientResponse); + requests.forEach(ProxyRequest::completeWithUnknownError); + return; + } + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ROUTER_OUT],[RESPONSE],response={}", clientResponse); + } + AutomqZoneRouterResponse zoneRouterResponse = (AutomqZoneRouterResponse) clientResponse.responseBody(); + handleRouterResponse(zoneRouterResponse, requests); + }).exceptionally(ex -> { + LOGGER.error("[ROUTER_OUT],[REQUEST_FAIL]", ex); + requests.forEach(ProxyRequest::completeWithUnknownError); + return null; + }); + } + + private void handleRouterResponse(AutomqZoneRouterResponse zoneRouterResponse, + List proxyRequests) { + List responses = zoneRouterResponse.data().responses(); + for (int i = 0; i < proxyRequests.size(); i++) { + ProxyRequest proxyRequest = proxyRequests.get(i); + AutomqZoneRouterResponseData.Response response = responses.get(i); + ProduceResponseData produceResponseData = ZoneRouterResponseCodec.decode(Unpooled.wrappedBuffer(response.data())); + response.setData(null); // gc the data + Map rst = new HashMap<>(); + produceResponseData.responses().forEach(topicData -> { + topicData.partitionResponses().forEach(partitionData -> { + ProduceResponse.PartitionResponse partitionResponse = new ProduceResponse.PartitionResponse( + Errors.forCode(partitionData.errorCode()), + partitionData.baseOffset(), + 0, // last offset , the network layer don't need + partitionData.logAppendTimeMs(), + partitionData.logStartOffset(), + partitionData.recordErrors().stream().map(e -> new ProduceResponse.RecordError(e.batchIndex(), e.batchIndexErrorMessage())).collect(Collectors.toList()), + partitionData.errorMessage(), + partitionData.currentLeader() + ); + rst.put(new TopicPartition(topicData.name(), partitionData.index()), partitionResponse); + }); + }); + proxyRequest.cf.complete(rst); + } + } + + /** + * Split the produce request to different nodes + */ + private Map split( + short apiVersion, + ClientIdMetadata clientId, + int timeout, + short requiredAcks, + short flag, + String transactionId, + Map entriesPerPartition + ) { + Map>> node2Entries = new HashMap<>(); + entriesPerPartition.forEach((tp, records) -> { + Node node = mapping.getRouteOutNode(tp.topic(), tp.partition(), clientId); + node2Entries.compute(node, (n, list) -> { + if (list == null) { + list = new ArrayList<>(); + } + list.add(Map.entry(tp, records)); + return list; + }); + }); + Map rst = new HashMap<>(); + node2Entries.forEach((node, entries) -> { + AtomicInteger size = new AtomicInteger(); + ProduceRequestData data = new ProduceRequestData(); + data.setTransactionalId(transactionId); + data.setAcks(requiredAcks); + data.setTimeoutMs(timeout); + + Map> topicData = new HashMap<>(); + entries.forEach(e -> { + TopicPartition tp = e.getKey(); + MemoryRecords records = e.getValue(); + topicData.compute(tp.topic(), (topicName, map) -> { + if (map == null) { + map = new HashMap<>(); + } + map.put(tp.partition(), records); + size.addAndGet(records.sizeInBytes()); + return map; + }); + }); + ProduceRequestData.TopicProduceDataCollection list = new ProduceRequestData.TopicProduceDataCollection(); + topicData.forEach((topicName, partitionData) -> { + list.add( + new ProduceRequestData.TopicProduceData() + .setName(topicName) + .setPartitionData( + partitionData.entrySet() + .stream() + .map(e -> new ProduceRequestData.PartitionProduceData().setIndex(e.getKey()).setRecords(e.getValue())) + .collect(Collectors.toList()) + ) + ); + }); + data.setTopicData(list); + rst.put(node, new ProxyRequest(apiVersion, flag, data, size.get())); + }); + return rst; + } + + private void cleanup() { + try { + List objects = objectStorage.list(ZoneRouterPack.getObjectPathPrefix(currentNode.id())).get(); + objectStorage.delete(objects.stream().map(o -> (ObjectStorage.ObjectPath) o).collect(Collectors.toList())).get(); + } catch (Throwable e) { + LOGGER.error("cleanup fail", e); + } + } + + static class ProxyRequest { + short apiVersion; + short flag; + ProduceRequestData data; + int size; + List topicPartitions; + CompletableFuture> cf; + + public ProxyRequest(short apiVersion, short flag, ProduceRequestData data, int size) { + this.apiVersion = apiVersion; + this.flag = flag; + this.data = data; + this.size = size; + this.cf = new CompletableFuture<>(); + this.topicPartitions = new ArrayList<>(); + this.data.topicData().forEach(topicData -> topicData.partitionData().forEach(partitionData -> { + topicPartitions.add(new TopicPartition(topicData.name(), partitionData.index())); + })); + } + + public void afterRouter() { + data = null; // gc the data + } + + public void completeWithNotLeaderNotFollower() { + completeWithError(Errors.NOT_LEADER_OR_FOLLOWER); + } + + public void completeWithUnknownError() { + completeWithError(Errors.UNKNOWN_SERVER_ERROR); + } + + private void completeWithError(Errors errors) { + Map rst = new HashMap<>(); + topicPartitions.forEach(tp -> rst.put(tp, new ProduceResponse.PartitionResponse(errors, -1, -1, -1, -1, Collections.emptyList(), ""))); + cf.complete(rst); + } + } + + static class PerfMode { + private final boolean enable = Systems.getEnvBool("AUTOMQ_PERF_MODE", false); // force route 2 / 3 traffic + private final int availableZoneCount = Systems.getEnvInt("AUTOMQ_PERF_MODE_AZ_COUNT", 3); + private final Map routerMap = new ConcurrentHashMap<>(); + private final AtomicInteger routerIndex = new AtomicInteger(); + + public boolean isRouteInPerfMode(Map entriesPerPartition) { + if (!enable) { + return false; + } + long producerId = entriesPerPartition.entrySet().iterator().next().getValue().batches().iterator().next().producerId(); + return routerMap.computeIfAbsent(producerId, k -> { + int index = routerIndex.incrementAndGet(); + return index % availableZoneCount != 0; + }); + } + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterOutV2.java b/core/src/main/java/kafka/automq/zerozone/RouterOutV2.java new file mode 100644 index 0000000000..5b44b49c27 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterOutV2.java @@ -0,0 +1,388 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ProduceRequestArgs; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.message.AutomqZoneRouterRequestData; +import org.apache.kafka.common.message.AutomqZoneRouterResponseData; +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.message.ProduceResponseData; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.requests.ProduceResponse; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterRequest; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; +import org.apache.kafka.common.utils.Time; + +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Queue; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class RouterOutV2 { + private static final Logger LOGGER = LoggerFactory.getLogger(RouterOutV2.class); + private final Node currentNode; + private final RouterChannel routerChannel; + private final Map proxies = new ConcurrentHashMap<>(); + + private final LocalProxy localProxy; + private final GetRouterOutNode mapping; + private final AsyncSender asyncSender; + private final Time time; + + public RouterOutV2(Node currentNode, RouterChannel routerChannel, GetRouterOutNode mapping, + NonBlockingLocalRouterHandler localRouterHandler, AsyncSender asyncSender, Time time) { + this.currentNode = currentNode; + this.routerChannel = routerChannel; + this.mapping = mapping; + this.localProxy = new LocalProxy(localRouterHandler); + this.asyncSender = asyncSender; + this.time = time; + } + + public void handleProduceAppendProxy(ProduceRequestArgs args) { + long timeoutMillis = time.milliseconds() + args.timeout(); + short flag = new ZoneRouterProduceRequest.Flag().internalTopicsAllowed(args.internalTopicsAllowed()).value(); + Map responseMap = new ConcurrentHashMap<>(); + List> cfList = new ArrayList<>(args.entriesPerPartition().size()); + long startNanos = time.nanoseconds(); + boolean acks0 = args.requiredAcks() == (short) 0; + for (Map.Entry entry : args.entriesPerPartition().entrySet()) { + TopicPartition tp = entry.getKey(); + MemoryRecords records = entry.getValue(); + Node node = mapping.getRouteOutNode(tp.topic(), tp.partition(), args.clientId()); + if (node.id() == Node.noNode().id()) { + responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.NOT_LEADER_OR_FOLLOWER)); + continue; + } + short orderHint = orderHint(tp, args.clientId().connectionId()); + int recordSize = records.sizeInBytes(); + ZoneRouterProduceRequest zoneRouterProduceRequest = zoneRouterProduceRequest(args, flag, tp, records); + CompletableFuture channelCf = routerChannel.append(node.id(), orderHint, ZoneRouterPackWriter.encodeDataBlock(List.of(zoneRouterProduceRequest))); + CompletableFuture proxyCf = channelCf.thenCompose(channelRst -> { + long timeNanos = time.nanoseconds(); + ZeroZoneMetricsManager.APPEND_CHANNEL_LATENCY.record(timeNanos - startNanos); + ProxyRequest proxyRequest = new ProxyRequest(tp, channelRst.epoch(), channelRst.channelOffset(), zoneRouterProduceRequest, recordSize, timeoutMillis); + sendProxyRequest(node, proxyRequest); + return proxyRequest.cf.thenAccept(response -> { + if (!acks0) { + responseMap.put(tp, response); + } + ZeroZoneMetricsManager.PROXY_REQUEST_LATENCY.record(time.nanoseconds() - startNanos); + }); + }).exceptionally(ex -> { + LOGGER.error("Exception in processing append proxies", ex); + // Make the producer retry send. + responseMap.put(tp, errorPartitionResponse(Errors.LEADER_NOT_AVAILABLE)); + return null; + }); + cfList.add(proxyCf); + } + Consumer> responseCallback = args.responseCallback(); + if (acks0) { + // When acks=0 is set, invoke the callback directly without waiting for data persistence to complete. + args.entriesPerPartition().forEach((tp, records) -> + responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.NONE))); + responseCallback.accept(responseMap); + } else { + CompletableFuture cf = CompletableFuture.allOf(cfList.toArray(new CompletableFuture[0])); + cf.thenAccept(nil -> responseCallback.accept(responseMap)).exceptionally(ex -> { + LOGGER.error("[UNEXPECTED],[ROUTE_FAIL]", ex); + return null; + }); + } + } + + private static short orderHint(TopicPartition tp, String connectionId) { + return (short) (Objects.hash(tp.topic(), tp.partition(), connectionId) % Short.MAX_VALUE); + } + + private void sendProxyRequest(Node node, ProxyRequest proxyRequest) { + if (node.id() == currentNode.id()) { + localProxy.send(proxyRequest); + } else { + Proxy proxy = proxies.computeIfAbsent(node, RemoteProxy::new); + proxy.send(proxyRequest); + } + } + + interface Proxy { + void send(ProxyRequest request); + } + + static ProduceResponse.PartitionResponse errorPartitionResponse(Errors error) { + return new ProduceResponse.PartitionResponse(error, -1, -1, -1, -1, Collections.emptyList(), ""); + } + + static class LocalProxy implements Proxy { + private final NonBlockingLocalRouterHandler localRouterHandler; + + LocalProxy(NonBlockingLocalRouterHandler handler) { + localRouterHandler = handler; + } + + @Override + public void send(ProxyRequest request) { + localRouterHandler.append(ChannelOffset.of(request.channelOffset), request.zoneRouterProduceRequest) + .whenComplete((resp, ex) -> { + if (ex != null) { + request.completeWithError(Errors.forException(ex)); + } else { + ProduceResponseData produceResponseData = ZoneRouterResponseCodec.decode(Unpooled.wrappedBuffer(resp.data())); + resp.setData(null); // gc the data + ProduceResponseData.PartitionProduceResponse partitionData = produceResponseData.responses().iterator().next() + .partitionResponses().get(0); + request.cf.complete(partitionResponse(partitionData)); + } + }); + } + } + + class RemoteProxy implements Proxy { + private static final int MAX_INFLIGHT_SIZE = 64; + private static final long LINGER_NANOS = TimeUnit.MICROSECONDS.toNanos(100); + private final Node node; + private final Semaphore inflightLimiter = new Semaphore(MAX_INFLIGHT_SIZE); + private final Queue requestBatchQueue = new ConcurrentLinkedQueue<>(); + private RequestBatch requestBatch = null; + + public RemoteProxy(Node node) { + this.node = node; + } + + public synchronized void send(ProxyRequest request) { + ZeroZoneMetricsManager.recordRouterOutBytes(node.id(), request.recordSize); + synchronized (this) { + if (requestBatch == null) { + requestBatch = new RequestBatch(time, LINGER_NANOS, 8192); + Threads.COMMON_SCHEDULER.schedule(() -> trySendRequestBatch(requestBatch), LINGER_NANOS, TimeUnit.NANOSECONDS); + } + if (requestBatch.add(request)) { + requestBatchQueue.add(requestBatch); + requestBatch = null; + trySendRequestBatch(null); + } + } + } + + private synchronized void trySendRequestBatch(RequestBatch forceSend) { + if (inflightLimiter.availablePermits() == 0) { + return; + } + if ((requestBatch != null && requestBatch.lingerTimeout()) + || (forceSend != null && requestBatch == forceSend)) { + requestBatchQueue.add(requestBatch); + requestBatch = null; + } + for (; ; ) { + RequestBatch waitingSend = requestBatchQueue.peek(); + if (waitingSend == null) { + break; + } + if (!inflightLimiter.tryAcquire()) { + break; + } + requestBatchQueue.poll(); + List> futures = new ArrayList<>(); + waitingSend.getRequests().forEach((epoch, requests) -> futures.add(batchSend(epoch, requests))); + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])) + .whenComplete((nil, ex) -> { + synchronized (RemoteProxy.this) { + inflightLimiter.release(); + trySendRequestBatch(null); + } + }); + } + } + + private CompletableFuture batchSend(long epoch, List requests) { + RouterRecordV2 routerRecord = new RouterRecordV2( + currentNode.id(), + requests.stream().map(r -> r.channelOffset).collect(Collectors.toList()) + ); + AutomqZoneRouterRequest.Builder builder = new AutomqZoneRouterRequest.Builder( + new AutomqZoneRouterRequestData().setMetadata(routerRecord.encode().array()) + .setVersion((short) 1) + .setRouteEpoch(epoch) + ); + return asyncSender.sendRequest(node, builder).thenAccept(clientResponse -> { + if (!clientResponse.hasResponse()) { + LOGGER.error("[ROUTER_OUT],[NO_RESPONSE],node={},response={}", node, clientResponse); + // Make the producer retry send. + requests.forEach(r -> r.completeWithError(Errors.LEADER_NOT_AVAILABLE)); + return; + } + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ROUTER_OUT],[RESPONSE],node={},response={}", node, clientResponse); + } + AutomqZoneRouterResponse zoneRouterResponse = (AutomqZoneRouterResponse) clientResponse.responseBody(); + handleRouterResponse(zoneRouterResponse, requests); + }).exceptionally(ex -> { + LOGGER.error("[ROUTER_OUT],[REQUEST_FAIL],node={}", node, ex); + requests.forEach(r -> r.completeWithError(Errors.LEADER_NOT_AVAILABLE)); + return null; + }); + } + + private void handleRouterResponse(AutomqZoneRouterResponse zoneRouterResponse, List requests) { + short errorCode = zoneRouterResponse.data().errorCode(); + if (errorCode == Errors.UNKNOWN_SERVER_ERROR.code()) { + // We could find the detail error log in the rpc server side. + // Set the error to LEADER_NOT_AVAILABLE to make the producer retry sending. + errorCode = Errors.LEADER_NOT_AVAILABLE.code(); + } + if (errorCode != Errors.NONE.code()) { + Errors error = Errors.forCode(errorCode); + requests.forEach(r -> r.completeWithError(error)); + return; + } + List responses = zoneRouterResponse.data().responses(); + Iterator requestIt = requests.iterator(); + Iterator responseIt = responses.iterator(); + while (requestIt.hasNext()) { + ProxyRequest request = requestIt.next(); + AutomqZoneRouterResponseData.Response response = responseIt.next(); + ProduceResponseData produceResponseData = ZoneRouterResponseCodec.decode(Unpooled.wrappedBuffer(response.data())); + response.setData(null); // gc the data + ProduceResponseData.PartitionProduceResponse partitionData = produceResponseData.responses().iterator().next() + .partitionResponses().get(0); + request.cf.complete(partitionResponse(partitionData)); + } + + } + } + + static class RequestBatch { + private final long batchStartNanos; + + private final Time time; + private final long lingerNanos; + private final long batchSize; + private final Map> requests = new TreeMap<>(); + + public RequestBatch(Time time, long lingerNanos, int batchSize) { + this.time = time; + this.lingerNanos = lingerNanos; + this.batchSize = batchSize; + this.batchStartNanos = System.nanoTime(); + } + + /** + * Add request to batch + * + * @param request {@link ProxyRequest} + * @return whether the batch is full + */ + public boolean add(ProxyRequest request) { + requests.computeIfAbsent(request.epoch, key -> new ArrayList<>()).add(request); + return requests.size() > batchSize || time.nanoseconds() - batchStartNanos >= lingerNanos; + } + + public boolean lingerTimeout() { + return time.nanoseconds() - batchStartNanos >= lingerNanos; + } + + public Map> getRequests() { + return requests; + } + } + + static class ProxyRequest { + final TopicPartition topicPartition; + final long epoch; + final ByteBuf channelOffset; + final ZoneRouterProduceRequest zoneRouterProduceRequest; + final int recordSize; + final long timeoutMillis; + final CompletableFuture cf = new CompletableFuture<>(); + + public ProxyRequest(TopicPartition topicPartition, long epoch, ByteBuf channelOffset, + ZoneRouterProduceRequest zoneRouterProduceRequest, int recordSize, long timeoutMillis) { + this.topicPartition = topicPartition; + this.epoch = epoch; + this.channelOffset = channelOffset; + this.zoneRouterProduceRequest = zoneRouterProduceRequest; + this.recordSize = recordSize; + this.timeoutMillis = timeoutMillis; + } + + public void completeWithUnknownError() { + completeWithError(Errors.UNKNOWN_SERVER_ERROR); + } + + private void completeWithError(Errors errors) { + ProduceResponse.PartitionResponse rst = errorPartitionResponse(errors); + cf.complete(rst); + } + } + + private static ZoneRouterProduceRequest zoneRouterProduceRequest(ProduceRequestArgs args, short flag, + TopicPartition tp, + MemoryRecords records) { + ProduceRequestData data = new ProduceRequestData(); + data.setTransactionalId(args.transactionId()); + data.setAcks(args.requiredAcks()); + data.setTimeoutMs(args.timeout()); + ProduceRequestData.TopicProduceDataCollection topics = new ProduceRequestData.TopicProduceDataCollection(); + ProduceRequestData.TopicProduceData topic = new ProduceRequestData.TopicProduceData(); + topic.setName(tp.topic()); + topic.setPartitionData(List.of(new ProduceRequestData.PartitionProduceData().setIndex(tp.partition()).setRecords(records))); + topics.add(topic); + data.setTopicData(topics); + return new ZoneRouterProduceRequest(args.apiVersion(), flag, data); + } + + private static ProduceResponse.PartitionResponse partitionResponse( + ProduceResponseData.PartitionProduceResponse partitionData) { + return new ProduceResponse.PartitionResponse( + Errors.forCode(partitionData.errorCode()), + partitionData.baseOffset(), + 0, // last offset , the network layer don't need + partitionData.logAppendTimeMs(), + partitionData.logStartOffset(), + partitionData.recordErrors().stream().map(e -> new ProduceResponse.RecordError(e.batchIndex(), e.batchIndexErrorMessage())).collect(Collectors.toList()), + partitionData.errorMessage(), + partitionData.currentLeader() + ); + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterRecord.java b/core/src/main/java/kafka/automq/zerozone/RouterRecord.java new file mode 100644 index 0000000000..ff52f7e8c9 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterRecord.java @@ -0,0 +1,97 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class RouterRecord { + private static final short MAGIC = 0x01; + private final int nodeId; + private final short bucketId; + private final long objectId; + private final int position; + private final int size; + + public RouterRecord(int nodeId, short bucketId, long objectId, int position, int size) { + this.nodeId = nodeId; + this.bucketId = bucketId; + this.objectId = objectId; + this.position = position; + this.size = size; + } + + public int nodeId() { + return nodeId; + } + + public short bucketId() { + return bucketId; + } + + public long objectId() { + return objectId; + } + + public int position() { + return position; + } + + public int size() { + return size; + } + + public ByteBuf encode() { + ByteBuf buf = Unpooled.buffer(1 /* magic */ + 4 /* nodeId */ + 2 /* bucketId */ + 8 /* objectId */ + 4 /* position */ + 4 /* size */); + buf.writeByte(MAGIC); + buf.writeInt(nodeId); + buf.writeShort(bucketId); + buf.writeLong(objectId); + buf.writeInt(position); + buf.writeInt(size); + return buf; + } + + @Override + public String toString() { + return "RouterRecord{" + + "nodeId=" + nodeId + + ", bucketId=" + bucketId + + ", objectId=" + objectId + + ", position=" + position + + ", size=" + size + + '}'; + } + + public static RouterRecord decode(ByteBuf buf) { + buf = buf.slice(); + byte magic = buf.readByte(); + if (magic != MAGIC) { + throw new IllegalArgumentException("Invalid magic byte: " + magic); + } + int nodeId = buf.readInt(); + short bucketId = buf.readShort(); + long objectId = buf.readLong(); + int position = buf.readInt(); + int size = buf.readInt(); + return new RouterRecord(nodeId, bucketId, objectId, position, size); + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/RouterRecordV2.java b/core/src/main/java/kafka/automq/zerozone/RouterRecordV2.java new file mode 100644 index 0000000000..42b0397dad --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/RouterRecordV2.java @@ -0,0 +1,75 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import java.util.ArrayList; +import java.util.List; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class RouterRecordV2 { + private static final short MAGIC = 0x02; + private final int nodeId; + private final List channelOffsets; + + public RouterRecordV2(int nodeId, List channelOffsets) { + this.nodeId = nodeId; + this.channelOffsets = channelOffsets; + } + + public int nodeId() { + return nodeId; + } + + public List channelOffsets() { + return channelOffsets; + } + + public ByteBuf encode() { + int size = 1 /* magic */ + 4 /* nodeId */ + channelOffsets.stream().mapToInt(buf -> buf.readableBytes() + 2).sum(); + ByteBuf buf = Unpooled.buffer(size); + buf.writeByte(MAGIC); + buf.writeInt(nodeId); + channelOffsets.forEach(channelOffset -> { + buf.writeShort(channelOffset.readableBytes()); + buf.writeBytes(channelOffset.duplicate()); + }); + return buf; + } + + public static RouterRecordV2 decode(ByteBuf buf) { + buf = buf.slice(); + byte magic = buf.readByte(); + if (magic != MAGIC) { + throw new IllegalArgumentException("Invalid magic byte: " + magic); + } + int nodeId = buf.readInt(); + List channelOffsets = new ArrayList<>(buf.readableBytes() / 16); + while (buf.readableBytes() > 0) { + short size = buf.readShort(); + ByteBuf channelOffset = Unpooled.buffer(size); + buf.readBytes(channelOffset); + channelOffsets.add(channelOffset); + } + return new RouterRecordV2(nodeId, channelOffsets); + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/SnapshotReadPartitionsManager.java b/core/src/main/java/kafka/automq/zerozone/SnapshotReadPartitionsManager.java new file mode 100644 index 0000000000..9b0fcfc750 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/SnapshotReadPartitionsManager.java @@ -0,0 +1,549 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.cluster.Partition; +import kafka.cluster.PartitionSnapshot; +import kafka.log.streamaspect.LazyStream; +import kafka.server.KafkaConfig; +import kafka.server.MetadataCache; +import kafka.server.streamaspect.ElasticReplicaManager; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicIdPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.image.MetadataDelta; +import org.apache.kafka.image.MetadataImage; +import org.apache.kafka.image.loader.MetadataListener; +import org.apache.kafka.metadata.BrokerRegistration; +import org.apache.kafka.server.common.automq.AutoMQVersion; + +import com.automq.stream.s3.cache.SnapshotReadCache; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.threads.EventLoop; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalLong; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; + +public class SnapshotReadPartitionsManager implements MetadataListener, ProxyTopologyChangeListener { + private static final Logger LOGGER = LoggerFactory.getLogger(SnapshotReadPartitionsManager.class); + static final long REQUEST_INTERVAL_MS = 1; + private final KafkaConfig config; + private final Time time; + private final ConfirmWALProvider confirmWALProvider; + private final ElasticReplicaManager replicaManager; + private final MetadataCache metadataCache; + private final AsyncSender asyncSender; + private final Replayer replayer; + private final ScheduledExecutorService scheduler = Threads.newSingleThreadScheduledExecutor("AUTOMQ_SNAPSHOT_READ", true, LOGGER); + private final Map topicId2name = new ConcurrentHashMap<>(); + private final CacheEventListener cacheEventListener = new CacheEventListener(); + final Map subscribers = new HashMap<>(); + // all snapshot read partition changes exec in a single eventloop to ensure the thread-safe. + final EventLoop eventLoop = new EventLoop("AUTOMQ_SNAPSHOT_READ_WORKER"); + private AutoMQVersion version; + private volatile boolean closed = false; + private final List> closingSubscribers = new CopyOnWriteArrayList<>(); + + public SnapshotReadPartitionsManager(KafkaConfig config, Metrics metrics, Time time, + ConfirmWALProvider confirmWALProvider, + ElasticReplicaManager replicaManager, MetadataCache metadataCache, Replayer replayer) { + this.config = config; + this.time = time; + this.confirmWALProvider = confirmWALProvider; + this.replicaManager = replicaManager; + this.metadataCache = metadataCache; + this.replayer = replayer; + this.asyncSender = new AsyncSender.BrokersAsyncSender(config, metrics, "snapshot_read", Time.SYSTEM, "AUTOMQ_SNAPSHOT_READ", new LogContext()); + } + + // test only + SnapshotReadPartitionsManager(KafkaConfig config, Time time, ConfirmWALProvider confirmWALProvider, + ElasticReplicaManager replicaManager, + MetadataCache metadataCache, Replayer replayer, AsyncSender asyncSender) { + this.config = config; + this.time = time; + this.confirmWALProvider = confirmWALProvider; + this.replicaManager = replicaManager; + this.metadataCache = metadataCache; + this.replayer = replayer; + this.asyncSender = asyncSender; + } + + public synchronized void close() { + closed = true; + subscribers.forEach((k, s) -> s.close()); + CompletableFuture.allOf(closingSubscribers.toArray(new CompletableFuture[0])).join(); + subscribers.clear(); + } + + @Override + public void onChange(MetadataDelta delta, MetadataImage image) { + if (delta.topicsDelta() != null && !delta.topicsDelta().deletedTopicIds().isEmpty()) { + Set deletedTopicIds = delta.topicsDelta().deletedTopicIds(); + scheduler.schedule(() -> deletedTopicIds.forEach(topicId2name::remove), 1, TimeUnit.MINUTES); + } + triggerSubscribersApply(); + } + + public void setVersion(AutoMQVersion newVersion) { + AutoMQVersion oldVersion = this.version; + this.version = newVersion; + if (oldVersion != null && (oldVersion.isZeroZoneV2Supported() != newVersion.isZeroZoneV2Supported())) { + // reset the subscriber + resetSubscribers(newVersion); + } + } + + public synchronized CompletableFuture nextSnapshotCf() { + return CompletableFuture.allOf(subscribers.values().stream() + .map(Subscriber::nextSnapshotCf) + .toList() + .toArray(new CompletableFuture[0]) + ); + } + + private synchronized void triggerSubscribersApply() { + subscribers.forEach((nodeId, subscriber) -> subscriber.apply()); + } + + private synchronized void resetSubscribers(AutoMQVersion version) { + Set nodes = subscribers.keySet(); + nodes.forEach(nodeId -> subscribers.computeIfPresent(nodeId, (id, subscribe) -> { + subscribe.close(); + return new Subscriber(subscribe.node, version); + })); + } + + private void removePartition(TopicIdPartition topicIdPartition, Partition expected) { + replicaManager.computeSnapshotReadPartition(topicIdPartition.topicPartition(), (tp, current) -> { + if (current == null || expected == current) { + expected.close(); + LOGGER.info("[SNAPSHOT_READ_REMOVE],tp={},epoch={}", topicIdPartition, expected.getLeaderEpoch()); + return null; + } + // The expected partition was closed when the current partition put in. + return current; + }); + } + + private Optional addPartition(TopicIdPartition topicIdPartition, PartitionSnapshot snapshot) { + AtomicReference ref = new AtomicReference<>(); + Supplier newPartition = () -> { + Partition partition = replicaManager.newSnapshotReadPartition(topicIdPartition); + partition.snapshot(snapshot); + ref.set(partition); + LOGGER.info("[SNAPSHOT_READ_ADD],tp={},epoch={}", topicIdPartition, ref.get().getLeaderEpoch()); + return partition; + }; + replicaManager.computeSnapshotReadPartition(topicIdPartition.topicPartition(), (tp, current) -> { + if (current == null) { + return newPartition.get(); + } + if (!topicIdPartition.topicId().equals(current.topicId().get())) { + if (metadataCache.getTopicName(topicIdPartition.topicId()).isDefined()) { + LOGGER.warn("[SNAPSHOT_READ_ADD],[KICK_OUT_DELETED_TOPIC],tp={},snapshot={},current={}", topicIdPartition, snapshot, current.topicId()); + current.close(); + return newPartition.get(); + } else { + LOGGER.warn("[SNAPSHOT_READ_ADD],[IGNORE],tp={},snapshot={},", topicIdPartition, snapshot); + return current; + } + } + if (snapshot.leaderEpoch() > current.getLeaderEpoch()) { + // The partition is reassigned from N1 to N2. + // Both N1 and N2 are subscribed by the current node. + // The N2 ADD operation is arrived before N1 REMOVE operation. + LOGGER.warn("[SNAPSHOT_READ_ADD],[REMOVE_OLD_EPOCH],tp={},snapshot={},currentEpoch={}", topicIdPartition, snapshot, current.getLeaderEpoch()); + current.close(); + return newPartition.get(); + } else { + LOGGER.warn("[SNAPSHOT_READ_ADD],[OLD_EPOCH],tp={},snapshot={},currentEpoch={}", topicIdPartition, snapshot, current.getLeaderEpoch()); + } + return current; + }); + return Optional.ofNullable(ref.get()); + } + + @Override + public synchronized void onChange(Map> main2proxyByRack) { + if (closed) { + return; + } + Set newSubscribeNodes = calSubscribeNodes(main2proxyByRack, config.nodeId()); + subscribers.entrySet().removeIf(entry -> { + if (!newSubscribeNodes.contains(entry.getKey())) { + entry.getValue().close(); + return true; + } else { + return false; + } + }); + newSubscribeNodes.forEach(nodeId -> { + if (!subscribers.containsKey(nodeId)) { + Optional opt = metadataCache.getNode(nodeId).node(config.interBrokerListenerName().value()); + if (opt.isPresent()) { + subscribers.put(nodeId, new Subscriber(opt.get(), version)); + } else { + LOGGER.error("[SNAPSHOT_READ_SUBSCRIBE],node={} not found", nodeId); + } + } + }); + } + + public SnapshotReadCache.EventListener cacheEventListener() { + return cacheEventListener; + } + + private String getTopicName(Uuid topicId) { + return topicId2name.computeIfAbsent(topicId, id -> metadataCache.topicIdsToNames().get(id)); + } + + // only for test + Subscriber newSubscriber(Node node, AutoMQVersion version, SubscriberRequester requester, + SubscriberReplayer dataLoader) { + return new Subscriber(node, version, requester, dataLoader); + } + + class Subscriber { + final Node node; + final Map partitions = new HashMap<>(); + boolean closed; + long appliedCount = 0; + final AtomicLong applyingCount = new AtomicLong(); + final Queue waitingMetadataReadyQueue = new LinkedList<>(); + final Queue waitingDataLoadedQueue = new LinkedList<>(); + final Queue snapshotWithOperations = new LinkedList<>(); + private final SubscriberRequester requester; + private final SubscriberReplayer replayer; + private final AutoMQVersion version; + + public Subscriber(Node node, AutoMQVersion version) { + this.node = node; + this.version = version; + this.replayer = new SubscriberReplayer(confirmWALProvider, SnapshotReadPartitionsManager.this.replayer, node, metadataCache); + this.requester = new SubscriberRequester(this, node, version, asyncSender, SnapshotReadPartitionsManager.this::getTopicName, eventLoop, time); + // start the tasks after initialized. + this.requester.start(); + run(); + LOGGER.info("[SNAPSHOT_READ_SUBSCRIBE],node={}", node); + } + + // only for test + public Subscriber(Node node, AutoMQVersion version, SubscriberRequester requester, + SubscriberReplayer replayer) { + this.node = node; + this.version = version; + this.requester = requester; + this.replayer = replayer; + } + + public void apply() { + applyingCount.incrementAndGet(); + eventLoop.execute(() -> { + long applyingCount = this.applyingCount.get(); + if (this.appliedCount == applyingCount) { + return; + } + unsafeRun(); + this.appliedCount = applyingCount; + }); + } + + public void requestCommit() { + if (version.isZeroZoneV2Supported()) { + eventLoop.execute(() -> requester.requestCommit = true); + } + } + + /** + * Get the next snapshot future. The future will be completed after the next sync snapshots have been applied. + */ + public CompletableFuture nextSnapshotCf() { + return requester.nextSnapshotCf(); + } + + public CompletableFuture close() { + LOGGER.info("[SNAPSHOT_READ_UNSUBSCRIBE],node={}", node); + CompletableFuture cf = new CompletableFuture<>(); + closingSubscribers.add(cf); + cf.whenComplete((nil, ex) -> closingSubscribers.remove(cf)); + eventLoop.execute(() -> { + try { + closed = true; + requester.close(); + partitions.forEach(SnapshotReadPartitionsManager.this::removePartition); + partitions.clear(); + snapshotWithOperations.clear(); + CompletableFuture replayerCloseCf = replayer.close(); + requester.nextSnapshotCf().complete(null); + FutureUtil.propagate(replayerCloseCf, cf); + } catch (Throwable e) { + cf.completeExceptionally(e); + } + }); + return cf; + } + + void run() { + eventLoop.execute(this::unsafeRun); + } + + /** + * Must run in eventLoop. + */ + void unsafeRun() { + try { + this.run0(); + } catch (Throwable e) { + LOGGER.error("[SNAPSHOT_SUBSCRIBE_ERROR]", e); + reset("SUBSCRIBE_ERROR: " + e.getMessage()); + scheduler.schedule(this::run, 1, TimeUnit.SECONDS); + } + } + + private void run0() { + if (closed) { + return; + } + // try replay the SSO/WAL data to snapshot-read cache + tryReplay(); + // after the metadata is ready and data is preload in snapshot-read cache, + // then apply the snapshot to partition. + applySnapshot(); + } + + void reset(String reason) { + LOGGER.info("[SNAPSHOT_READ_SUBSCRIBER_RESET],node={},reason={}", node, reason); + partitions.forEach(SnapshotReadPartitionsManager.this::removePartition); + partitions.clear(); + waitingMetadataReadyQueue.clear(); + snapshotWithOperations.clear(); + waitingDataLoadedQueue.clear(); + requester.reset(); + replayer.reset(); + } + + void onNewWalEndOffset(String walConfig, RecordOffset endOffset, byte[] walDeltaData) { + replayer.onNewWalEndOffset(walConfig, endOffset, walDeltaData); + } + + void onNewOperationBatch(OperationBatch batch) { + waitingMetadataReadyQueue.add(batch); + } + + void applySnapshot() { + while (!snapshotWithOperations.isEmpty()) { + SnapshotWithOperation snapshotWithOperation = snapshotWithOperations.peek(); + if (snapshotWithOperation.isSnapshotMark()) { + snapshotWithOperations.poll(); + snapshotWithOperation.snapshotCf.complete(null); + continue; + } + TopicIdPartition topicIdPartition = snapshotWithOperation.topicIdPartition; + + switch (snapshotWithOperation.operation) { + case ADD: { + Optional partition = addPartition(topicIdPartition, snapshotWithOperation.snapshot); + if (partition.isEmpty()) { + reset(String.format("Cannot find partition %s", topicIdPartition)); + return; + } + partition.ifPresent(p -> partitions.put(topicIdPartition, p)); + snapshotWithOperations.poll(); + break; + } + case PATCH: { + Partition partition = partitions.get(topicIdPartition); + if (partition != null) { + partition.snapshot(snapshotWithOperation.snapshot); + } else { + LOGGER.error("[SNAPSHOT_READ_PATCH],[SKIP],{}", snapshotWithOperation); + } + snapshotWithOperations.poll(); + break; + } + case REMOVE: { + Partition partition = partitions.remove(topicIdPartition); + if (partition != null) { + removePartition(topicIdPartition, partition); + } + snapshotWithOperations.poll(); + break; + } + default: + throw new IllegalArgumentException("SnapshotOperation " + snapshotWithOperation.operation + " is not supported"); + } + } + } + + void tryReplay() { + // - ZERO_ZONE_V0: Collect all the operation which data metadata is ready in kraft and replay the SSO. + // - ZERO_ZONE_V1: Directly replay the WAL. + List batches = new ArrayList<>(); + for (; ; ) { + OperationBatch batch = waitingMetadataReadyQueue.peek(); + if (batch == null) { + break; + } + if (version.isZeroZoneV2Supported() || checkBatchMetadataReady0(batch, metadataCache)) { + waitingMetadataReadyQueue.poll(); + batches.add(batch); + } else { + break; + } + } + if (batches.isEmpty()) { + return; + } + + CompletableFuture waitingDataLoadedCf; + if (version.isZeroZoneV2Supported()) { + waitingDataLoadedCf = replayer.replayWal(); + } else { + // Trigger incremental SSO data loading. + waitingDataLoadedCf = replayer.relayObject(); + } + WaitingDataLoadTask task = new WaitingDataLoadTask(time.milliseconds(), batches, waitingDataLoadedCf); + waitingDataLoadedQueue.add(task); + // After the SSO data loads to the snapshot-read cache, then apply operations. + waitingDataLoadedCf.thenAcceptAsync(nil -> checkDataLoaded(), eventLoop); + } + + void checkDataLoaded() { + for (; ; ) { + WaitingDataLoadTask task = waitingDataLoadedQueue.peek(); + if (task == null || !task.cf.isDone()) { + break; + } + task.operationBatchList.forEach(batch -> snapshotWithOperations.addAll(batch.operations)); + waitingDataLoadedQueue.poll(); + unsafeRun(); + } + } + } + + class CacheEventListener implements SnapshotReadCache.EventListener { + @Override + public void onEvent(SnapshotReadCache.Event event) { + synchronized (SnapshotReadPartitionsManager.this) { + if (event instanceof SnapshotReadCache.RequestCommitEvent) { + Subscriber subscriber = subscribers.get(((SnapshotReadCache.RequestCommitEvent) event).nodeId()); + if (subscriber != null) { + subscriber.requestCommit(); + } + } + } + } + } + + static boolean checkBatchMetadataReady0(OperationBatch batch, MetadataCache metadataCache) { + for (; ; ) { + if (batch.readyIndex == batch.operations.size() - 1) { + return true; + } + SnapshotWithOperation operation = batch.operations.get(batch.readyIndex + 1); + if (!operation.isSnapshotMark() && isMetadataUnready(operation.snapshot.streamEndOffsets(), metadataCache)) { + return false; + } + batch.readyIndex = batch.readyIndex + 1; + } + } + + static boolean isMetadataUnready(Map streamEndOffsets, MetadataCache metadataCache) { + AtomicBoolean ready = new AtomicBoolean(true); + streamEndOffsets.forEach((streamId, endOffset) -> { + if (streamId == LazyStream.NOOP_STREAM_ID) { + return; + } + OptionalLong opt = metadataCache.getStreamEndOffset(streamId); + if (opt.isEmpty()) { + throw new RuntimeException(String.format("Cannot find streamId=%s, the kraft metadata replay delay or the topic is deleted.", streamId)); + } + long endOffsetInKraft = opt.getAsLong(); + if (endOffsetInKraft < endOffset) { + ready.set(false); + } + }); + return !ready.get(); + } + + static Set calSubscribeNodes(Map> main2proxyByRack, + int currentNodeId) { + Set nodes = new HashSet<>(); + main2proxyByRack.forEach((rack, main2proxy) -> { + main2proxy.forEach((mainNodeId, proxy) -> { + if (proxy.id() == currentNodeId) { + nodes.add(mainNodeId); + } + }); + }); + return nodes; + } + + static class OperationBatch { + final List operations; + int readyIndex; + + public OperationBatch() { + this.operations = new ArrayList<>(); + this.readyIndex = -1; + } + } + + static class WaitingDataLoadTask { + final long timestamp; + final List operationBatchList; + final CompletableFuture cf; + + public WaitingDataLoadTask(long timestamp, List operationBatchList, + CompletableFuture cf) { + this.timestamp = timestamp; + this.operationBatchList = operationBatchList; + this.cf = cf; + } + } + +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/zerozone/SnapshotWithOperation.java b/core/src/main/java/kafka/automq/zerozone/SnapshotWithOperation.java new file mode 100644 index 0000000000..428892506e --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/SnapshotWithOperation.java @@ -0,0 +1,64 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.partition.snapshot.SnapshotOperation; +import kafka.cluster.PartitionSnapshot; + +import org.apache.kafka.common.TopicIdPartition; + +import java.util.concurrent.CompletableFuture; + +public class SnapshotWithOperation { + final TopicIdPartition topicIdPartition; + final PartitionSnapshot snapshot; + final SnapshotOperation operation; + final CompletableFuture snapshotCf; + + public SnapshotWithOperation(TopicIdPartition topicIdPartition, PartitionSnapshot snapshot, + SnapshotOperation operation) { + this(topicIdPartition, snapshot, operation, null); + } + + public SnapshotWithOperation(TopicIdPartition topicIdPartition, PartitionSnapshot snapshot, + SnapshotOperation operation, CompletableFuture snapshotCf) { + this.topicIdPartition = topicIdPartition; + this.snapshot = snapshot; + this.operation = operation; + this.snapshotCf = snapshotCf; + } + + public static SnapshotWithOperation snapshotMark(CompletableFuture cf) { + return new SnapshotWithOperation(null, null, null, cf); + } + + public boolean isSnapshotMark() { + return snapshotCf != null; + } + + @Override + public String toString() { + return "SnapshotWithOperation{" + + "topicIdPartition=" + topicIdPartition + + ", snapshot=" + snapshot + + ", operation=" + operation + + '}'; + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/SubscriberReplayer.java b/core/src/main/java/kafka/automq/zerozone/SubscriberReplayer.java new file mode 100644 index 0000000000..d9e9f2a20f --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/SubscriberReplayer.java @@ -0,0 +1,166 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.server.MetadataCache; + +import org.apache.kafka.common.Node; +import org.apache.kafka.image.S3ObjectsImage; +import org.apache.kafka.metadata.stream.S3Object; +import org.apache.kafka.metadata.stream.S3StreamSetObject; + +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.WriteAheadLog; +import com.automq.stream.utils.FutureUtil; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.function.LongConsumer; +import java.util.stream.Collectors; + +import static kafka.automq.partition.snapshot.ConfirmWalDataDelta.decodeDeltaRecords; + +class SubscriberReplayer { + private static final Logger LOGGER = LoggerFactory.getLogger(SubscriberReplayer.class); + private static final ExecutorService CLOSE_EXECUTOR = Executors.newCachedThreadPool(); + private long loadedObjectOrderId = -1L; + private CompletableFuture lastDataLoadCf = CompletableFuture.completedFuture(null); + private CompletableFuture wal; + private RecordOffset loadedEndOffset = null; + + private final Replayer replayer; + private final Node node; + private final MetadataCache metadataCache; + private final ConfirmWALProvider confirmWALProvider; + + public SubscriberReplayer(ConfirmWALProvider confirmWALProvider, Replayer replayer, Node node, + MetadataCache metadataCache) { + this.confirmWALProvider = confirmWALProvider; + this.replayer = replayer; + this.node = node; + this.metadataCache = metadataCache; + } + + public void onNewWalEndOffset(String walConfig, RecordOffset endOffset, byte[] walDeltaData) { + if (wal == null) { + this.wal = confirmWALProvider.readOnly(walConfig, node.id()); + } + if (endOffset.equals(loadedEndOffset)) { + return; + } + RecordOffset startOffset = this.loadedEndOffset; + this.loadedEndOffset = endOffset; + if (startOffset == null) { + return; + } + // The replayer will ensure the order of replay + this.lastDataLoadCf = wal.thenCompose(w -> replayer.replay(w, startOffset, endOffset, decodeDeltaRecords(walDeltaData)).thenAccept(nil -> { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("replay {} confirm wal [{}, {})", node, startOffset, endOffset); + } + })).exceptionally(ex -> { + LOGGER.error("[UNEXPECTED] replay confirm wal fail", ex); + return null; + }); + } + + public CompletableFuture relayObject() { + List newObjects = nextObjects().stream().filter(object -> { + if (object.objectSize() > 200L * 1024 * 1024) { + LOGGER.warn("The object {} is bigger than 200MiB, skip load it", object); + return false; + } else { + return true; + } + }).collect(Collectors.toList()); + if (newObjects.isEmpty()) { + return lastDataLoadCf; + } + long loadedObjectOrderId = this.loadedObjectOrderId; + return lastDataLoadCf = lastDataLoadCf.thenCompose(nil -> replayer.replay(newObjects)).thenAccept(nil -> { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[LOAD_SNAPSHOT_READ_DATA],node={},loadedObjectOrderId={},newObjects={}", node, loadedObjectOrderId, newObjects); + } + }); + } + + public CompletableFuture replayWal() { + return lastDataLoadCf; + } + + public CompletableFuture close() { + CompletableFuture wal = this.wal; + if (wal != null) { + return CompletableFuture.runAsync(() -> FutureUtil.suppress(() -> wal.get().shutdownGracefully(), LOGGER), CLOSE_EXECUTOR); + } + return CompletableFuture.completedFuture(null); + } + + public void reset() { + this.loadedObjectOrderId = -1L; + this.loadedEndOffset = null; + } + + private List nextObjects() { + return nextObjects0(metadataCache, node.id(), loadedObjectOrderId, value -> loadedObjectOrderId = value); + } + + static List nextObjects0(MetadataCache metadataCache, int nodeId, long loadedObjectOrderId, + LongConsumer loadedObjectOrderIdUpdater) { + return metadataCache.safeRun(image -> { + List newObjects = new ArrayList<>(); + List streamSetObjects = image.streamsMetadata().getStreamSetObjects(nodeId); + S3ObjectsImage objectsImage = image.objectsMetadata(); + long nextObjectOrderId = loadedObjectOrderId; + if (loadedObjectOrderId == -1L) { + // try to load the latest 16MB data + long size = 0; + for (int i = streamSetObjects.size() - 1; i >= 0 && size < 16 * 1024 * 1024 && newObjects.size() < 8; i--) { + S3StreamSetObject sso = streamSetObjects.get(i); + S3Object s3object = objectsImage.getObjectMetadata(sso.objectId()); + size += s3object.getObjectSize(); + newObjects.add(new S3ObjectMetadata(sso.objectId(), s3object.getObjectSize(), s3object.getAttributes())); + nextObjectOrderId = Math.max(nextObjectOrderId, sso.orderId()); + } + } else { + for (int i = streamSetObjects.size() - 1; i >= 0; i--) { + S3StreamSetObject sso = streamSetObjects.get(i); + if (sso.orderId() <= loadedObjectOrderId) { + break; + } + S3Object s3object = objectsImage.getObjectMetadata(sso.objectId()); + newObjects.add(new S3ObjectMetadata(sso.objectId(), s3object.getObjectSize(), s3object.getAttributes())); + nextObjectOrderId = Math.max(nextObjectOrderId, sso.orderId()); + } + } + loadedObjectOrderIdUpdater.accept(nextObjectOrderId); + Collections.reverse(newObjects); + return newObjects; + }); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/SubscriberRequester.java b/core/src/main/java/kafka/automq/zerozone/SubscriberRequester.java new file mode 100644 index 0000000000..9d0b59bb7b --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/SubscriberRequester.java @@ -0,0 +1,290 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.partition.snapshot.SnapshotOperation; +import kafka.cluster.PartitionSnapshot; +import kafka.log.streamaspect.ElasticLogMeta; +import kafka.log.streamaspect.ElasticStreamSegmentMeta; +import kafka.log.streamaspect.SliceRange; + +import org.apache.kafka.clients.ClientResponse; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicIdPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotRequestData; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.s3.AutomqGetPartitionSnapshotRequest; +import org.apache.kafka.common.requests.s3.AutomqGetPartitionSnapshotResponse; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.server.common.automq.AutoMQVersion; +import org.apache.kafka.storage.internals.log.LogOffsetMetadata; +import org.apache.kafka.storage.internals.log.TimestampOffset; + +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.threads.EventLoop; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; +import java.util.stream.Collectors; + +import io.netty.buffer.Unpooled; + +@SuppressWarnings("NPathComplexity") class SubscriberRequester { + private static final Logger LOGGER = LoggerFactory.getLogger(SubscriberRequester.class); + private boolean closed = false; + private long lastRequestTime; + private int sessionId; + private int sessionEpoch; + boolean requestCommit = false; + boolean requestReset = false; + private CompletableFuture nextSnapshotCf = new CompletableFuture<>(); + + private final SnapshotReadPartitionsManager.Subscriber subscriber; + private final Node node; + private final AutoMQVersion version; + private final AsyncSender asyncSender; + private final Function topicNameGetter; + private final EventLoop eventLoop; + private final Time time; + + public SubscriberRequester(SnapshotReadPartitionsManager.Subscriber subscriber, Node node, AutoMQVersion version, + AsyncSender asyncSender, + Function topicNameGetter, EventLoop eventLoop, Time time) { + this.subscriber = subscriber; + this.node = node; + this.version = version; + this.asyncSender = asyncSender; + this.topicNameGetter = topicNameGetter; + this.eventLoop = eventLoop; + this.time = time; + } + + public void start() { + request(); + } + + public void reset() { + requestReset = true; + } + + public void close() { + closed = true; + } + + public CompletableFuture nextSnapshotCf() { + return nextSnapshotCf; + } + + private void request() { + eventLoop.execute(this::request0); + } + + private void request0() { + if (closed) { + return; + } + // The snapshotCf will be completed after all snapshots in the response have been applied. + CompletableFuture snapshotCf = this.nextSnapshotCf; + this.nextSnapshotCf = new CompletableFuture<>(); + // The request may fail. So when the nextSnapshotCf complete, we will complete the current snapshotCf. + FutureUtil.propagate(nextSnapshotCf, snapshotCf); + + tryReset0(); + lastRequestTime = time.milliseconds(); + AutomqGetPartitionSnapshotRequestData data = new AutomqGetPartitionSnapshotRequestData().setSessionId(sessionId).setSessionEpoch(sessionEpoch); + if (version.isZeroZoneV2Supported()) { + data.setVersion((short) 1); + } + if (version.isZeroZoneV2Supported() && sessionEpoch == 0) { + // request ConfirmWAL commit data to main storage, then the data that doesn't replay could be read from main storage. + data.setRequestCommit(true); + } else if (requestCommit) { + requestCommit = false; + data.setRequestCommit(true); + } + if (data.requestCommit()) { + LOGGER.info("[SNAPSHOT_SUBSCRIBE_REQUEST_COMMIT],node={},sessionId={},sessionEpoch={}", node, sessionId, sessionEpoch); + } + AutomqGetPartitionSnapshotRequest.Builder builder = new AutomqGetPartitionSnapshotRequest.Builder(data); + asyncSender.sendRequest(node, builder) + .thenAcceptAsync(rst -> { + try { + handleResponse(rst, snapshotCf); + } catch (Exception e) { + subscriber.reset("Exception when handle snapshot response: " + e.getMessage()); + } + subscriber.unsafeRun(); + }, eventLoop) + .exceptionally(ex -> { + LOGGER.error("[SNAPSHOT_SUBSCRIBE_ERROR],node={}", node, ex); + return null; + }).whenComplete((nil, ex) -> { + long elapsed = time.milliseconds() - lastRequestTime; + if (SnapshotReadPartitionsManager.REQUEST_INTERVAL_MS > elapsed) { + Threads.COMMON_SCHEDULER.schedule(() -> eventLoop.execute(this::request), SnapshotReadPartitionsManager.REQUEST_INTERVAL_MS - elapsed, TimeUnit.MILLISECONDS); + } else { + request(); + } + }); + } + + private void handleResponse(ClientResponse clientResponse, CompletableFuture snapshotCf) { + if (closed) { + return; + } + if (tryReset0()) { + // If it needs to reset, then drop the response. + return; + } + if (!clientResponse.hasResponse()) { + if (clientResponse.wasDisconnected() || clientResponse.wasTimedOut()) { + LOGGER.warn("[GET_SNAPSHOTS],[REQUEST_FAIL],response={}", clientResponse); + } else { + LOGGER.error("[GET_SNAPSHOTS],[NO_RESPONSE],response={}", clientResponse); + } + return; + } + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[GET_SNAPSHOTS],[RESPONSE],response={}", clientResponse); + } + AutomqGetPartitionSnapshotResponse zoneRouterResponse = (AutomqGetPartitionSnapshotResponse) clientResponse.responseBody(); + AutomqGetPartitionSnapshotResponseData resp = zoneRouterResponse.data(); + if (resp.errorCode() != Errors.NONE.code()) { + LOGGER.error("[GET_SNAPSHOTS],[ERROR],response={}", resp); + return; + } + if (sessionId != 0 && resp.sessionId() != sessionId) { + // switch to a new session + subscriber.reset(String.format("switch sessionId from %s to %s", sessionId, resp.sessionId())); + // reset immediately to the new session. + tryReset0(); + } + sessionId = resp.sessionId(); + sessionEpoch = resp.sessionEpoch(); + SnapshotReadPartitionsManager.OperationBatch batch = new SnapshotReadPartitionsManager.OperationBatch(); + resp.topics().forEach(topic -> topic.partitions().forEach(partition -> { + String topicName = topicNameGetter.apply(topic.topicId()); + if (topicName == null) { + String reason = String.format("Cannot find topic uuid=%s, the kraft metadata replay delay or the topic is deleted.", topic.topicId()); + subscriber.reset(reason); + throw new RuntimeException(reason); + } + batch.operations.add(convert(new TopicIdPartition(topic.topicId(), partition.partitionIndex(), topicName), partition)); + })); + // Make sure the REMOVE operations will be applied first. + batch.operations.sort((o1, o2) -> { + int c1 = o1.operation.code() == SnapshotOperation.REMOVE.code() ? 0 : 1; + int c2 = o2.operation.code() == SnapshotOperation.REMOVE.code() ? 0 : 1; + return c1 - c2; + }); + short requestVersion = clientResponse.requestHeader().apiVersion(); + if (resp.confirmWalEndOffset() != null && resp.confirmWalEndOffset().length > 0) { + // zerozone v2 + subscriber.onNewWalEndOffset(resp.confirmWalConfig(), + DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset())), + requestVersion >= 2 ? resp.confirmWalDeltaData() : null); + } + batch.operations.add(SnapshotWithOperation.snapshotMark(snapshotCf)); + subscriber.onNewOperationBatch(batch); + } + + private boolean tryReset0() { + if (requestReset) { + sessionId = 0; + sessionEpoch = 0; + requestReset = false; + return true; + } else { + return false; + } + } + + static SnapshotWithOperation convert(TopicIdPartition topicIdPartition, + AutomqGetPartitionSnapshotResponseData.PartitionSnapshot src) { + PartitionSnapshot.Builder snapshot = PartitionSnapshot.builder(); + snapshot.leaderEpoch(src.leaderEpoch()); + snapshot.logMeta(convert(src.logMetadata())); + snapshot.firstUnstableOffset(convert(src.firstUnstableOffset())); + snapshot.logEndOffset(convert(src.logEndOffset())); + src.streamMetadata().forEach(m -> snapshot.streamEndOffset(m.streamId(), m.endOffset())); + snapshot.lastTimestampOffset(convertTimestampOffset(src.lastTimestampOffset())); + + SnapshotOperation operation = SnapshotOperation.parse(src.operation()); + return new SnapshotWithOperation(topicIdPartition, snapshot.build(), operation); + } + + static ElasticLogMeta convert(AutomqGetPartitionSnapshotResponseData.LogMetadata src) { + if (src == null || src.segments().isEmpty()) { + // the AutomqGetPartitionSnapshotResponseData's default LogMetadata is an empty LogMetadata. + return null; + } + ElasticLogMeta logMeta = new ElasticLogMeta(); + logMeta.setStreamMap(src.streamMap().stream().collect(Collectors.toMap(AutomqGetPartitionSnapshotResponseData.StreamMapping::name, AutomqGetPartitionSnapshotResponseData.StreamMapping::streamId))); + src.segments().forEach(m -> logMeta.getSegmentMetas().add(convert(m))); + return logMeta; + } + + static ElasticStreamSegmentMeta convert(AutomqGetPartitionSnapshotResponseData.SegmentMetadata src) { + ElasticStreamSegmentMeta meta = new ElasticStreamSegmentMeta(); + meta.baseOffset(src.baseOffset()); + meta.createTimestamp(src.createTimestamp()); + meta.lastModifiedTimestamp(src.lastModifiedTimestamp()); + meta.streamSuffix(src.streamSuffix()); + meta.logSize(src.logSize()); + meta.log(convert(src.log())); + meta.time(convert(src.time())); + meta.txn(convert(src.transaction())); + meta.firstBatchTimestamp(src.firstBatchTimestamp()); + meta.timeIndexLastEntry(convert(src.timeIndexLastEntry())); + return meta; + } + + static SliceRange convert(AutomqGetPartitionSnapshotResponseData.SliceRange src) { + return SliceRange.of(src.start(), src.end()); + } + + static ElasticStreamSegmentMeta.TimestampOffsetData convert( + AutomqGetPartitionSnapshotResponseData.TimestampOffsetData src) { + return ElasticStreamSegmentMeta.TimestampOffsetData.of(src.timestamp(), src.offset()); + } + + static TimestampOffset convertTimestampOffset(AutomqGetPartitionSnapshotResponseData.TimestampOffsetData src) { + if (src == null) { + return null; + } + return new TimestampOffset(src.timestamp(), src.offset()); + } + + static LogOffsetMetadata convert(AutomqGetPartitionSnapshotResponseData.LogOffsetMetadata src) { + if (src == null) { + return null; + } + // The segment offset should be fill in Partition#snapshot + return new LogOffsetMetadata(src.messageOffset(), -1, src.relativePositionInSegment()); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZeroZoneMetricsManager.java b/core/src/main/java/kafka/automq/zerozone/ZeroZoneMetricsManager.java new file mode 100644 index 0000000000..20de21504a --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZeroZoneMetricsManager.java @@ -0,0 +1,72 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.s3.metrics.Metrics; +import com.automq.stream.s3.metrics.MetricsLevel; +import com.automq.stream.s3.metrics.wrapper.DeltaHistogram; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; + +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.LongCounter; + +public class ZeroZoneMetricsManager { + private static final String PREFIX = "kafka_zonerouter_"; + + private static final Cache ROUTER_OUT_ATTRIBUTES_CACHE = CacheBuilder.newBuilder() + .maximumSize(1000) + .expireAfterAccess(1, TimeUnit.MINUTES) + .build(); + private static final Cache ROUTER_IN_ATTRIBUTES_CACHE = CacheBuilder.newBuilder() + .maximumSize(1000) + .expireAfterAccess(1, TimeUnit.MINUTES) + .build(); + private static final LongCounter ROUTER_BYTES = Metrics.instance().counter(meter -> meter + .counterBuilder(PREFIX + "router_bytes") + .setUnit("bytes") + .setDescription("Cross zone router bytes") + .build()); + + private static final Metrics.HistogramBundle ROUTER_LATENCY = Metrics.instance().histogram(PREFIX + "router_latency", "ZeroZone route latency", "nanoseconds"); + public static final DeltaHistogram APPEND_CHANNEL_LATENCY = ROUTER_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "out", AttributeKey.stringKey("stage"), "append_channel")); + public static final DeltaHistogram PROXY_REQUEST_LATENCY = ROUTER_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "out", AttributeKey.stringKey("stage"), "proxy_request")); + public static final DeltaHistogram GET_CHANNEL_LATENCY = ROUTER_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "in", AttributeKey.stringKey("stage"), "get_channel")); + + public static void recordRouterOutBytes(int toNodeId, int bytes) { + try { + ROUTER_BYTES.add(bytes, ROUTER_OUT_ATTRIBUTES_CACHE.get(toNodeId, () -> Attributes.of(AttributeKey.stringKey("type"), "out", AttributeKey.stringKey("peerNodeId"), Integer.toString(toNodeId)))); + } catch (ExecutionException e) { + // suppress + } + } + + public static void recordRouterInBytes(int fromNodeId, int bytes) { + try { + ROUTER_BYTES.add(bytes, ROUTER_IN_ATTRIBUTES_CACHE.get(fromNodeId, () -> Attributes.of(AttributeKey.stringKey("type"), "in", AttributeKey.stringKey("peerNodeId"), Integer.toString(fromNodeId)))); + } catch (ExecutionException e) { + // suppress + } + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZeroZoneThreadLocalContext.java b/core/src/main/java/kafka/automq/zerozone/ZeroZoneThreadLocalContext.java new file mode 100644 index 0000000000..1455a87454 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZeroZoneThreadLocalContext.java @@ -0,0 +1,61 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import io.netty.util.concurrent.FastThreadLocal; + +public class ZeroZoneThreadLocalContext { + + private static final FastThreadLocal WRITE_CONTEXT = new FastThreadLocal<>() { + @Override protected WriteContext initialValue() { + return new WriteContext(); + } + }; + + public static WriteContext writeContext() { + return WRITE_CONTEXT.get(); + } + + public static void attach(WriteContext writeContext) { + WRITE_CONTEXT.set(writeContext); + } + + + public static class WriteContext { + ChannelOffset channelOffset; + + private WriteContext() {} + + public ChannelOffset channelOffset() { + return channelOffset; + } + + public void reset() { + channelOffset = null; + } + + public WriteContext detach() { + WRITE_CONTEXT.set(new WriteContext()); + return this; + } + } + + +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZeroZoneTrafficInterceptor.java b/core/src/main/java/kafka/automq/zerozone/ZeroZoneTrafficInterceptor.java new file mode 100644 index 0000000000..57d4ec1770 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZeroZoneTrafficInterceptor.java @@ -0,0 +1,268 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdKey; +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.automq.interceptor.ProduceRequestArgs; +import kafka.automq.interceptor.TrafficInterceptor; +import kafka.server.KafkaConfig; +import kafka.server.MetadataCache; +import kafka.server.streamaspect.ElasticKafkaApis; +import kafka.server.streamaspect.ElasticReplicaManager; + +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.message.AutomqZoneRouterRequestData; +import org.apache.kafka.common.message.MetadataResponseData; +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.requests.ProduceResponse; +import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.image.MetadataDelta; +import org.apache.kafka.image.MetadataImage; +import org.apache.kafka.image.loader.LoaderManifest; +import org.apache.kafka.image.publisher.MetadataPublisher; +import org.apache.kafka.server.common.automq.AutoMQVersion; + +import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; +import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +public class ZeroZoneTrafficInterceptor implements TrafficInterceptor, MetadataPublisher { + private static final Logger LOGGER = LoggerFactory.getLogger(ZeroZoneTrafficInterceptor.class); + private final ElasticKafkaApis kafkaApis; + private final ClientRackProvider clientRackProvider; + private final List config; + private final BucketURI bucketURI; + + private final ProxyNodeMapping mapping; + + private final RouterOut routerOut; + private final RouterIn routerIn; + + private final RouterChannelProvider routerChannelProvider; + private final RouterOutV2 routerOutV2; + private final RouterInV2 routerInV2; + private final CommittedEpochManager committedEpochManager; + + private final SnapshotReadPartitionsManager snapshotReadPartitionsManager; + private volatile AutoMQVersion version; + private final AtomicBoolean closed = new AtomicBoolean(false); + + public ZeroZoneTrafficInterceptor( + RouterChannelProvider routerChannelProvider, + ConfirmWALProvider confirmWALProvider, + ElasticKafkaApis kafkaApis, + MetadataCache metadataCache, + ClientRackProvider clientRackProvider, + KafkaConfig kafkaConfig) { + this.routerChannelProvider = routerChannelProvider; + this.kafkaApis = kafkaApis; + + if (kafkaConfig.rack().isEmpty()) { + throw new IllegalArgumentException("The node rack should be set when enable cross available zone router"); + } + + + String interBrokerListenerName = kafkaConfig.interBrokerListenerName().value(); + int nodeId = kafkaConfig.nodeId(); + Node currentNode = kafkaConfig.effectiveAdvertisedBrokerListeners() + .find(endpoint -> Objects.equals(interBrokerListenerName, endpoint.listenerName().value())) + .map(endpoint -> new Node(nodeId, endpoint.host(), endpoint.port())) + .get(); + + this.mapping = new ProxyNodeMapping(currentNode, kafkaConfig.rack().get(), interBrokerListenerName, metadataCache); + + Time time = Time.SYSTEM; + + AsyncSender asyncSender = new AsyncSender.BrokersAsyncSender(kafkaConfig, kafkaApis.metrics(), "zone_router", time, ZoneRouterPack.ZONE_ROUTER_CLIENT_ID, new LogContext()); + + this.config = kafkaConfig.automq().zoneRouterChannels().get(); + + //noinspection OptionalGetWithoutIsPresent + this.bucketURI = kafkaConfig.automq().zoneRouterChannels().get().get(0); + this.clientRackProvider = clientRackProvider; + ObjectStorage objectStorage = ObjectStorageFactory.instance().builder(bucketURI) + .readWriteIsolate(true) + .inboundLimiter(GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.INBOUND)) + .outboundLimiter(GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.OUTBOUND)) + .build(); + this.routerOut = new RouterOut(currentNode, bucketURI, objectStorage, mapping::getRouteOutNode, kafkaApis, asyncSender, time); + this.routerIn = new RouterIn(objectStorage, kafkaApis, kafkaConfig.rack().get()); + + // Zero Zone V2 + this.routerInV2 = new RouterInV2(routerChannelProvider, kafkaApis, kafkaConfig.rack().get(), time); + this.routerOutV2 = new RouterOutV2(currentNode, routerChannelProvider.channel(), mapping::getRouteOutNode, routerInV2, asyncSender, time); + this.committedEpochManager = new CommittedEpochManager(nodeId); + this.routerChannelProvider.addEpochListener(committedEpochManager); + DefaultReplayer replayer = new DefaultReplayer(); + + this.version = metadataCache.autoMQVersion(); + + this.snapshotReadPartitionsManager = new SnapshotReadPartitionsManager(kafkaConfig, kafkaApis.metrics(), time, confirmWALProvider, + (ElasticReplicaManager) kafkaApis.replicaManager(), kafkaApis.metadataCache(), replayer); + this.snapshotReadPartitionsManager.setVersion(version); + kafkaApis.setSnapshotAwaitReadyProvider(this.snapshotReadPartitionsManager::nextSnapshotCf); + replayer.setCacheEventListener(this.snapshotReadPartitionsManager.cacheEventListener()); + mapping.registerListener(snapshotReadPartitionsManager); + + + LOGGER.info("start zero zone traffic interceptor with config={}", bucketURI); + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + committedEpochManager.close(); + snapshotReadPartitionsManager.close(); + } + } + + @Override + public void handleProduceRequest(ProduceRequestArgs args) { + if (closed.get()) { + Map responseMap = new HashMap<>(args.entriesPerPartition().size()); + args.entriesPerPartition().forEach((tp, records) -> + responseMap.put(tp, new ProduceResponse.PartitionResponse(Errors.NOT_LEADER_OR_FOLLOWER))); + args.responseCallback().accept(responseMap); + return; + } + ClientIdMetadata clientId = args.clientId(); + fillRackIfMissing(clientId); + if (version.isZeroZoneV2Supported()) { + routerOutV2.handleProduceAppendProxy(args); + } else { + if (clientId.rack() != null) { + routerOut.handleProduceAppendProxy(args); + } else { + MismatchRecorder.instance().record(args.entriesPerPartition().entrySet().iterator().next().getKey().topic(), clientId); + // If the client rack isn't set, then try to handle the request in the current node. + kafkaApis.handleProduceAppendJavaCompatible(args); + } + } + } + + @Override + public CompletableFuture handleZoneRouterRequest(AutomqZoneRouterRequestData request) { + if (request.version() == 0) { + return routerIn.handleZoneRouterRequest(request.metadata()); + } else { + ReentrantReadWriteLock.ReadLock readLock = committedEpochManager.readLock(); + readLock.lock(); + AtomicLong inflight = committedEpochManager.epochInflight(request.routeEpoch()); + inflight.incrementAndGet(); + try { + return routerInV2.handleZoneRouterRequest(request) + .whenComplete((resp, ex) -> inflight.decrementAndGet()); + } finally { + readLock.unlock(); + } + } + } + + @Override + public List handleMetadataResponse(ClientIdMetadata clientId, + List topics) { + fillRackIfMissing(clientId); + return mapping.handleMetadataResponse(clientId, topics); + } + + @Override + public Optional getLeaderNode(int leaderId, ClientIdMetadata clientId, + String listenerName) { + fillRackIfMissing(clientId); + return mapping.getLeaderNode(leaderId, clientId, listenerName); + } + + @Override + public String name() { + return "ObjectCrossZoneProduceRouter"; + } + + @Override + public void onMetadataUpdate(MetadataDelta delta, MetadataImage newImage, LoaderManifest manifest) { + if (closed.get()) { + return; + } + try { + mapping.onChange(delta, newImage); + snapshotReadPartitionsManager.onChange(delta, newImage); + version = newImage.features().autoMQVersion(); + this.snapshotReadPartitionsManager.setVersion(version); + routerChannelProvider.onChange(delta, newImage); + } catch (Throwable e) { + LOGGER.error("Failed to handle metadata update", e); + } + } + + public void setRouterInProduceHandler(RouterInProduceHandler routerInProduceHandler) { + routerIn.setRouterInProduceHandler(routerInProduceHandler); + routerInV2.setRouterInProduceHandler(routerInProduceHandler); + } + + private AutoMQVersion version() { + return version; + } + + @Override + public String toString() { + return "ZeroZoneTrafficInterceptor{config=" + config + '}'; + } + + private void fillRackIfMissing(ClientIdMetadata clientId) { + if (clientId.rack() == null) { + String rack = clientRackProvider.rack(clientId); + if (rack != null) { + clientId.metadata(ClientIdKey.AVAILABILITY_ZONE, List.of(rack)); + } + } + } + + static Map produceRequestToMap(ProduceRequestData data) { + Map realEntriesPerPartition = new HashMap<>(); + data.topicData().forEach(topicData -> + topicData.partitionData().forEach(partitionData -> + realEntriesPerPartition.put( + new TopicPartition(topicData.name(), partitionData.index()), + (MemoryRecords) partitionData.records() + ))); + return realEntriesPerPartition; + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZoneRouterPack.java b/core/src/main/java/kafka/automq/zerozone/ZoneRouterPack.java new file mode 100644 index 0000000000..886266b1b8 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZoneRouterPack.java @@ -0,0 +1,42 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import com.automq.stream.s3.metadata.ObjectUtils; + +public class ZoneRouterPack { + public static final byte PRODUCE_DATA_BLOCK_MAGIC = 0x01; + public static final int FOOTER_SIZE = 48; + public static final long PACK_MAGIC = 0x88e241b785f4cff9L; + + public static final String ZONE_ROUTER_CLIENT_ID = "__automq_zr"; + + public static String genObjectPath(int nodeId, long objectId) { + return getObjectPathPrefixBuilder(nodeId).append(objectId).toString(); + } + + public static String getObjectPathPrefix(int nodeId) { + return getObjectPathPrefixBuilder(nodeId).toString(); + } + + private static StringBuilder getObjectPathPrefixBuilder(int nodeId) { + return new StringBuilder(String.format("%08x", nodeId)).reverse().append("/").append(ObjectUtils.getNamespace()).append("/router/"); + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZoneRouterPackReader.java b/core/src/main/java/kafka/automq/zerozone/ZoneRouterPackReader.java new file mode 100644 index 0000000000..db62a4f5c4 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZoneRouterPackReader.java @@ -0,0 +1,82 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.protocol.ByteBufferAccessor; + +import com.automq.stream.s3.operator.ObjectStorage; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import io.netty.buffer.ByteBuf; + +import static kafka.automq.zerozone.ZoneRouterPack.PRODUCE_DATA_BLOCK_MAGIC; +import static kafka.automq.zerozone.ZoneRouterPack.genObjectPath; + +public class ZoneRouterPackReader { + private final short bucketId; + private final String path; + private final ObjectStorage objectStorage; + + public ZoneRouterPackReader(int nodeId, short bucketId, long objectId, ObjectStorage objectStorage) { + this.path = genObjectPath(nodeId, objectId); + this.bucketId = bucketId; + this.objectStorage = objectStorage; + } + + public CompletableFuture> readProduceRequests(Position position) { + return objectStorage + .rangeRead(new ObjectStorage.ReadOptions().bucket(bucketId), path, position.position(), position.position() + position.size()) + .thenApply(buf -> { + try { + return ZoneRouterPackReader.decodeDataBlock(buf); + } finally { + buf.release(); + } + }); + } + + /** + * Caution: ProduceRequestData$PartitionProduceData.records is a slice of buf. + */ + static List decodeDataBlock(ByteBuf buf) { + byte magic = buf.readByte(); + if (magic != PRODUCE_DATA_BLOCK_MAGIC) { + throw new IllegalArgumentException("Invalid magic byte: " + magic); + } + List requests = new ArrayList<>(); + while (buf.readableBytes() > 0) { + short apiVersion = buf.readShort(); + short flag = buf.readShort(); + int dataSize = buf.readInt(); + ByteBuf dataBuf = buf.slice(buf.readerIndex(), dataSize); + ProduceRequestData produceRequestData = new ProduceRequestData(); + produceRequestData.read(new ByteBufferAccessor(dataBuf.nioBuffer()), apiVersion); + buf.skipBytes(dataSize); + buf.retain(); + requests.add(new ZoneRouterProduceRequest(apiVersion, flag, produceRequestData, buf::release)); + } + return requests; + } + +} \ No newline at end of file diff --git a/core/src/main/java/kafka/automq/zerozone/ZoneRouterPackWriter.java b/core/src/main/java/kafka/automq/zerozone/ZoneRouterPackWriter.java new file mode 100644 index 0000000000..bf4dc102c6 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZoneRouterPackWriter.java @@ -0,0 +1,109 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.ObjectSerializationCache; + +import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.Writer; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; + +import static kafka.automq.zerozone.ZoneRouterPack.FOOTER_SIZE; +import static kafka.automq.zerozone.ZoneRouterPack.PACK_MAGIC; +import static kafka.automq.zerozone.ZoneRouterPack.genObjectPath; + +public class ZoneRouterPackWriter { + private final String path; + private final Writer writer; + private final CompositeByteBuf dataBuf; + + public ZoneRouterPackWriter(int nodeId, long objectId, ObjectStorage objectStorage) { + this.path = genObjectPath(nodeId, objectId); + this.dataBuf = ByteBufAlloc.compositeByteBuffer(); + this.writer = objectStorage.writer(new ObjectStorage.WriteOptions().enableFastRetry(true), path); + } + + public Position addProduceRequests(List produceRequests) { + int position = dataBuf.writerIndex(); + ByteBuf buf = encodeDataBlock(produceRequests); + int size = buf.readableBytes(); + dataBuf.addComponent(true, buf); + return new Position(position, size); + } + + public short bucketId() { + return writer.bucketId(); + } + + public ObjectStorage.ObjectPath objectPath() { + return new ObjectStorage.ObjectPath(bucketId(), path); + } + + public CompletableFuture close() { + ByteBuf footer = ByteBufAlloc.byteBuffer(FOOTER_SIZE); + footer.writeZero(40); + footer.writeLong(PACK_MAGIC); + dataBuf.addComponent(true, footer); + writer.write(dataBuf); + return writer.close(); + } + + public static ByteBuf encodeDataBlock(List produceRequests) { + int size = 1 /* magic */; + List objectSerializationCaches = new ArrayList<>(produceRequests.size()); + List dataSizes = new ArrayList<>(produceRequests.size()); + for (ZoneRouterProduceRequest produceRequest : produceRequests) { + + size += 2 /* api version */ + 2 /* flag */ + 4 /* data size */; + + ProduceRequestData data = produceRequest.data(); + ObjectSerializationCache objectSerializationCache = new ObjectSerializationCache(); + objectSerializationCaches.add(objectSerializationCache); + int dataSize = data.size(objectSerializationCache, produceRequest.apiVersion()); + dataSizes.add(dataSize); + size += dataSize; + } + ByteBuf buf = ByteBufAlloc.byteBuffer(size); + buf.writeByte(ZoneRouterPack.PRODUCE_DATA_BLOCK_MAGIC); + for (int i = 0; i < produceRequests.size(); i++) { + ZoneRouterProduceRequest produceRequest = produceRequests.get(i); + int dataSize = dataSizes.get(i); + ProduceRequestData data = produceRequest.data(); + ObjectSerializationCache objectSerializationCache = objectSerializationCaches.get(i); + + buf.writeShort(produceRequest.apiVersion()); + buf.writeShort(produceRequest.flag()); + buf.writeInt(dataSize); + data.write(new ByteBufferAccessor(buf.nioBuffer(buf.writerIndex(), dataSize)), objectSerializationCache, produceRequest.apiVersion()); + buf.writerIndex(buf.writerIndex() + dataSize); + } + return buf; + } + +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZoneRouterProduceRequest.java b/core/src/main/java/kafka/automq/zerozone/ZoneRouterProduceRequest.java new file mode 100644 index 0000000000..9850396a89 --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZoneRouterProduceRequest.java @@ -0,0 +1,120 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.message.ProduceRequestData; + +import java.util.Objects; + +import io.netty.util.AbstractReferenceCounted; +import io.netty.util.ReferenceCounted; + +public class ZoneRouterProduceRequest extends AbstractReferenceCounted implements AutoCloseable { + private final short apiVersion; + private final short flag; + private final ProduceRequestData data; + private final Runnable releaseHook; + + public ZoneRouterProduceRequest(short apiVersion, short flag, ProduceRequestData data) { + this(apiVersion, flag, data, () -> { + }); + } + + public ZoneRouterProduceRequest(short apiVersion, short flag, ProduceRequestData data, Runnable releaseHook) { + this.apiVersion = apiVersion; + this.data = data; + this.flag = flag; + this.releaseHook = releaseHook; + } + + public short apiVersion() { + return apiVersion; + } + + public short flag() { + return flag; + } + + public ProduceRequestData data() { + return data; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (o == null || getClass() != o.getClass()) + return false; + ZoneRouterProduceRequest request = (ZoneRouterProduceRequest) o; + return apiVersion == request.apiVersion && Objects.equals(data, request.data); + } + + @Override + public int hashCode() { + return Objects.hash(apiVersion, data); + } + + @Override + protected void deallocate() { + releaseHook.run(); + } + + @Override + public ReferenceCounted touch(Object o) { + return this; + } + + @Override + public void close() { + release(); + } + + public static class Flag { + private static final short INTERNAL_TOPICS_ALLOWED = 1; + + private short flag; + + public Flag(short flag) { + this.flag = flag; + } + + public Flag() { + this((short) 0); + } + + public short value() { + return flag; + } + + public Flag internalTopicsAllowed(boolean internalTopicsAllowed) { + if (internalTopicsAllowed) { + flag = (short) (flag | INTERNAL_TOPICS_ALLOWED); + } else { + flag = (short) (flag & ~INTERNAL_TOPICS_ALLOWED); + } + return this; + } + + public boolean internalTopicsAllowed() { + return (flag & INTERNAL_TOPICS_ALLOWED) != 0; + } + + } +} diff --git a/core/src/main/java/kafka/automq/zerozone/ZoneRouterResponseCodec.java b/core/src/main/java/kafka/automq/zerozone/ZoneRouterResponseCodec.java new file mode 100644 index 0000000000..b1c49ff51f --- /dev/null +++ b/core/src/main/java/kafka/automq/zerozone/ZoneRouterResponseCodec.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.message.ProduceResponseData; +import org.apache.kafka.common.protocol.ByteBufferAccessor; +import org.apache.kafka.common.protocol.ObjectSerializationCache; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class ZoneRouterResponseCodec { + public static final byte PRODUCE_RESPONSE_BLOCK_MAGIC = 0x01; + + public static ByteBuf encode(ProduceResponseData produceResponseData) { + short version = 11; + ObjectSerializationCache objectSerializationCache = new ObjectSerializationCache(); + int size = produceResponseData.size(objectSerializationCache, version); + ByteBuf buf = Unpooled.buffer(1 /* magic */ + 2 /* version */ + size); + buf.writeByte(PRODUCE_RESPONSE_BLOCK_MAGIC); + buf.writeShort(version); + produceResponseData.write(new ByteBufferAccessor(buf.nioBuffer(buf.writerIndex(), size)), objectSerializationCache, version); + buf.writerIndex(buf.writerIndex() + size); + return buf; + } + + public static ProduceResponseData decode(ByteBuf buf) { + byte magic = buf.readByte(); + if (magic != PRODUCE_RESPONSE_BLOCK_MAGIC) { + throw new IllegalArgumentException("Invalid magic byte: " + magic); + } + short version = buf.readShort(); + ProduceResponseData produceResponseData = new ProduceResponseData(); + produceResponseData.read(new ByteBufferAccessor(buf.nioBuffer(buf.readerIndex(), buf.readableBytes())), version); + return produceResponseData; + } + +} diff --git a/core/src/main/java/kafka/automq/zonerouter/ClientIdKey.java b/core/src/main/java/kafka/automq/zonerouter/ClientIdKey.java deleted file mode 100644 index 39170e18f2..0000000000 --- a/core/src/main/java/kafka/automq/zonerouter/ClientIdKey.java +++ /dev/null @@ -1,17 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.automq.zonerouter; - -public class ClientIdKey { - public static final String AVAILABILITY_ZONE = "automq_az"; - public static final String CLIENT_TYPE = "automq_type"; -} diff --git a/core/src/main/java/kafka/automq/zonerouter/ClientIdMetadata.java b/core/src/main/java/kafka/automq/zonerouter/ClientIdMetadata.java deleted file mode 100644 index b9abd93fda..0000000000 --- a/core/src/main/java/kafka/automq/zonerouter/ClientIdMetadata.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.automq.zonerouter; - -import com.automq.stream.utils.URIUtils; - -import java.net.InetAddress; -import java.util.List; -import java.util.Map; - -public class ClientIdMetadata { - private final String clientId; - private final Map> metadata; - private final InetAddress clientAddress; - - private ClientIdMetadata(String clientId, InetAddress clientAddress) { - this.clientId = clientId; - this.metadata = URIUtils.splitQuery(clientId); - this.clientAddress = clientAddress; - } - - public static ClientIdMetadata of(String clientId) { - return new ClientIdMetadata(clientId, null); - } - - public static ClientIdMetadata of(String clientId, InetAddress clientAddress) { - return new ClientIdMetadata(clientId, clientAddress); - } - - public String rack() { - List list = metadata.get(ClientIdKey.AVAILABILITY_ZONE); - if (list == null || list.isEmpty()) { - return null; - } - return list.get(0); - } - - public ClientType clientType() { - List list = metadata.get(ClientIdKey.CLIENT_TYPE); - if (list == null || list.isEmpty()) { - return null; - } - return ClientType.parse(list.get(0)); - } - - public String clientId() { - return clientId; - } - - @Override - public String toString() { - if (clientAddress == null) { - return clientId; - } else { - return clientId + "/" + clientAddress.getHostAddress(); - } - } -} diff --git a/core/src/main/java/kafka/automq/zonerouter/ClientType.java b/core/src/main/java/kafka/automq/zonerouter/ClientType.java deleted file mode 100644 index 60ae2397ce..0000000000 --- a/core/src/main/java/kafka/automq/zonerouter/ClientType.java +++ /dev/null @@ -1,28 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.automq.zonerouter; - -public enum ClientType { - PRODUCER, CONSUMER; - - public static ClientType parse(String str) { - switch (str) { - case "producer": - return PRODUCER; - case "consumer": - return CONSUMER; - default: - return null; - } - } - -} diff --git a/core/src/main/java/kafka/automq/zonerouter/NoopProduceRouter.java b/core/src/main/java/kafka/automq/zonerouter/NoopProduceRouter.java deleted file mode 100644 index c573b7e919..0000000000 --- a/core/src/main/java/kafka/automq/zonerouter/NoopProduceRouter.java +++ /dev/null @@ -1,86 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.automq.zonerouter; - -import kafka.server.MetadataCache; -import kafka.server.streamaspect.ElasticKafkaApis; - -import org.apache.kafka.common.Node; -import org.apache.kafka.common.TopicPartition; -import org.apache.kafka.common.message.MetadataResponseData; -import org.apache.kafka.common.network.ListenerName; -import org.apache.kafka.common.record.MemoryRecords; -import org.apache.kafka.common.record.RecordValidationStats; -import org.apache.kafka.common.requests.ProduceResponse; -import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; - -import com.automq.stream.utils.FutureUtil; - -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import java.util.function.Consumer; - -public class NoopProduceRouter implements ProduceRouter { - private final ElasticKafkaApis kafkaApis; - private final MetadataCache metadataCache; - - public NoopProduceRouter(ElasticKafkaApis kafkaApis, MetadataCache metadataCache) { - this.kafkaApis = kafkaApis; - this.metadataCache = metadataCache; - } - - @Override - public void handleProduceRequest(short apiVersion, ClientIdMetadata clientId, int timeout, short requiredAcks, - boolean internalTopicsAllowed, String transactionId, Map entriesPerPartition, - Consumer> responseCallback, - Consumer> recordValidationStatsCallback) { - kafkaApis.handleProduceAppendJavaCompatible( - timeout, - requiredAcks, - internalTopicsAllowed, - transactionId, - entriesPerPartition, - rst -> { - responseCallback.accept(rst); - return null; - }, - rst -> { - recordValidationStatsCallback.accept(rst); - return null; - }, - apiVersion - ); - } - - @Override - public CompletableFuture handleZoneRouterRequest(byte[] metadata) { - return FutureUtil.failedFuture(new UnsupportedOperationException()); - } - - @Override - public List handleMetadataResponse(String clientId, - List topics) { - return topics; - } - - @Override - public Optional getLeaderNode(int leaderId, ClientIdMetadata clientId, String listenerName) { - scala.Option opt = metadataCache.getAliveBrokerNode(leaderId, new ListenerName(listenerName)); - if (opt.isEmpty()) { - return Optional.empty(); - } else { - return Optional.of(opt.get()); - } - } -} diff --git a/core/src/main/java/kafka/automq/zonerouter/ProduceRouter.java b/core/src/main/java/kafka/automq/zonerouter/ProduceRouter.java deleted file mode 100644 index 34fd1ba017..0000000000 --- a/core/src/main/java/kafka/automq/zonerouter/ProduceRouter.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.automq.zonerouter; - -import org.apache.kafka.common.Node; -import org.apache.kafka.common.TopicPartition; -import org.apache.kafka.common.message.MetadataResponseData; -import org.apache.kafka.common.record.MemoryRecords; -import org.apache.kafka.common.record.RecordValidationStats; -import org.apache.kafka.common.requests.ProduceResponse; -import org.apache.kafka.common.requests.s3.AutomqZoneRouterResponse; - -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import java.util.function.Consumer; - -public interface ProduceRouter { - - void handleProduceRequest( - short apiVersion, - ClientIdMetadata clientId, - int timeout, - short requiredAcks, - boolean internalTopicsAllowed, - String transactionId, - Map entriesPerPartition, - Consumer> responseCallback, - Consumer> recordValidationStatsCallback - ); - - CompletableFuture handleZoneRouterRequest(byte[] metadata); - - List handleMetadataResponse(String clientId, - List topics); - - Optional getLeaderNode(int leaderId, ClientIdMetadata clientId, String listenerName); - -} diff --git a/core/src/main/java/kafka/cluster/LogEventListener.java b/core/src/main/java/kafka/cluster/LogEventListener.java new file mode 100644 index 0000000000..30ba50e16b --- /dev/null +++ b/core/src/main/java/kafka/cluster/LogEventListener.java @@ -0,0 +1,34 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.cluster; + +import org.apache.kafka.storage.internals.log.LogSegment; + +public interface LogEventListener { + + void onChanged(LogSegment segment, Event event); + + enum Event { + SEGMENT_CREATE, + SEGMENT_DELETE, + SEGMENT_UPDATE + } + +} diff --git a/core/src/main/java/kafka/cluster/PartitionSnapshot.java b/core/src/main/java/kafka/cluster/PartitionSnapshot.java new file mode 100644 index 0000000000..f71e943a36 --- /dev/null +++ b/core/src/main/java/kafka/cluster/PartitionSnapshot.java @@ -0,0 +1,149 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.cluster; + +import kafka.log.streamaspect.ElasticLogMeta; + +import org.apache.kafka.storage.internals.log.LogOffsetMetadata; +import org.apache.kafka.storage.internals.log.TimestampOffset; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +public class PartitionSnapshot { + private final int leaderEpoch; + private final ElasticLogMeta logMeta; + private final LogOffsetMetadata firstUnstableOffset; + private final LogOffsetMetadata logEndOffset; + private final Map streamEndOffsets; + private final TimestampOffset lastTimestampOffset; + private final CompletableFuture completeCf; + + public PartitionSnapshot(int leaderEpoch, ElasticLogMeta meta, LogOffsetMetadata firstUnstableOffset, LogOffsetMetadata logEndOffset, + Map offsets, TimestampOffset lastTimestampOffset, CompletableFuture completeCf) { + this.leaderEpoch = leaderEpoch; + this.logMeta = meta; + this.firstUnstableOffset = firstUnstableOffset; + this.logEndOffset = logEndOffset; + this.streamEndOffsets = offsets; + this.lastTimestampOffset = lastTimestampOffset; + this.completeCf = completeCf; + } + + public int leaderEpoch() { + return leaderEpoch; + } + + public ElasticLogMeta logMeta() { + return logMeta; + } + + public LogOffsetMetadata firstUnstableOffset() { + return firstUnstableOffset; + } + + public LogOffsetMetadata logEndOffset() { + return logEndOffset; + } + + public Map streamEndOffsets() { + return streamEndOffsets; + } + + public TimestampOffset lastTimestampOffset() { + return lastTimestampOffset; + } + + public CompletableFuture completeCf() { + return completeCf; + } + + @Override + public String toString() { + return "PartitionSnapshot{" + + "leaderEpoch=" + leaderEpoch + + ", logMeta=" + logMeta + + ", firstUnstableOffset=" + firstUnstableOffset + + ", logEndOffset=" + logEndOffset + + ", streamEndOffsets=" + streamEndOffsets + + ", lastTimestampOffset=" + lastTimestampOffset + + '}'; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private static final int DEFAULT_STREAM_COUNT = 4; + private int leaderEpoch; + private ElasticLogMeta logMeta; + private LogOffsetMetadata firstUnstableOffset; + private LogOffsetMetadata logEndOffset; + private final Map streamEndOffsets = new HashMap<>(DEFAULT_STREAM_COUNT); + private final List> streamLastAppendFutures = new ArrayList<>(DEFAULT_STREAM_COUNT); + private TimestampOffset lastTimestampOffset; + + public Builder leaderEpoch(int leaderEpoch) { + this.leaderEpoch = leaderEpoch; + return this; + } + + public Builder logMeta(ElasticLogMeta meta) { + this.logMeta = meta; + return this; + } + + public Builder firstUnstableOffset(LogOffsetMetadata firstUnstableOffset) { + this.firstUnstableOffset = firstUnstableOffset; + return this; + } + + public Builder logEndOffset(LogOffsetMetadata logEndOffset) { + this.logEndOffset = logEndOffset; + return this; + } + + public Builder streamEndOffset(long streamId, long endOffset) { + streamEndOffsets.put(streamId, endOffset); + return this; + } + + public Builder lastTimestampOffset(TimestampOffset lastTimestampOffset) { + this.lastTimestampOffset = lastTimestampOffset; + return this; + } + + public Builder addStreamLastAppendFuture(CompletableFuture future) { + if (future != null) { + streamLastAppendFutures.add(future); + } + return this; + } + + public PartitionSnapshot build() { + CompletableFuture doneCf = CompletableFuture.allOf(streamLastAppendFutures.toArray(new CompletableFuture[0])).exceptionally(ex -> null); + return new PartitionSnapshot(leaderEpoch, logMeta, firstUnstableOffset, logEndOffset, streamEndOffsets, lastTimestampOffset, doneCf); + } + } +} diff --git a/core/src/main/java/kafka/server/TelemetrySupport.java b/core/src/main/java/kafka/server/TelemetrySupport.java new file mode 100644 index 0000000000..7338b67e62 --- /dev/null +++ b/core/src/main/java/kafka/server/TelemetrySupport.java @@ -0,0 +1,139 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.server; + +import kafka.automq.table.metric.TableTopicMetricsManager; + +import org.apache.kafka.common.config.types.Password; +import org.apache.kafka.server.ProcessRole; +import org.apache.kafka.server.metrics.KafkaYammerMetrics; +import org.apache.kafka.server.metrics.s3stream.S3StreamKafkaMetricsManager; + +import com.automq.opentelemetry.AutoMQTelemetryManager; +import com.automq.opentelemetry.exporter.MetricsExportConfig; +import com.automq.shell.AutoMQApplication; +import com.automq.stream.s3.metrics.Metrics; +import com.automq.stream.s3.metrics.MetricsConfig; +import com.automq.stream.s3.metrics.MetricsLevel; +import com.automq.stream.s3.metrics.S3StreamMetricsManager; + +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.Meter; +import scala.collection.immutable.Set; + +/** + * Helper used by the core module to bootstrap AutoMQ telemetry using the AutoMQTelemetryManager implement. + */ +public final class TelemetrySupport { + private static final Logger LOGGER = LoggerFactory.getLogger(TelemetrySupport.class); + private static final String COMMON_JMX_PATH = "/jmx/rules/common.yaml"; + private static final String BROKER_JMX_PATH = "/jmx/rules/broker.yaml"; + private static final String CONTROLLER_JMX_PATH = "/jmx/rules/controller.yaml"; + private static final String KAFKA_METRICS_PREFIX = "kafka_stream_"; + + private TelemetrySupport() { + // Utility class + } + + public static AutoMQTelemetryManager start(KafkaConfig config, String clusterId) { + AutoMQTelemetryManager telemetryManager = new AutoMQTelemetryManager( + config.automq().metricsExporterURI(), + clusterId, + String.valueOf(config.nodeId()), + AutoMQApplication.getBean(MetricsExportConfig.class) + ); + + telemetryManager.setJmxConfigPaths(buildJmxConfigPaths(config)); + telemetryManager.init(); + telemetryManager.startYammerMetricsReporter(KafkaYammerMetrics.defaultRegistry()); + initializeMetrics(telemetryManager, config); + + return telemetryManager; + } + + private static void initializeMetrics(AutoMQTelemetryManager manager, KafkaConfig config) { + S3StreamKafkaMetricsManager.setTruststoreCertsSupplier(() -> { + try { + Password password = config.getPassword("ssl.truststore.certificates"); + return password != null ? password.value() : null; + } catch (Exception e) { + LOGGER.error("Failed to obtain truststore certificates", e); + return null; + } + }); + + S3StreamKafkaMetricsManager.setCertChainSupplier(() -> { + try { + Password password = config.getPassword("ssl.keystore.certificate.chain"); + return password != null ? password.value() : null; + } catch (Exception e) { + LOGGER.error("Failed to obtain certificate chain", e); + return null; + } + }); + + Meter meter = manager.getMeter(); + MetricsLevel metricsLevel = parseMetricsLevel(config.s3MetricsLevel()); + long metricsIntervalMs = (long) config.s3ExporterReportIntervalMs(); + MetricsConfig metricsConfig = new MetricsConfig(metricsLevel, Attributes.empty(), metricsIntervalMs); + + Metrics.instance().setup(meter, metricsConfig); + S3StreamMetricsManager.configure(new MetricsConfig(metricsLevel, Attributes.empty(), metricsIntervalMs)); + S3StreamMetricsManager.initMetrics(meter, KAFKA_METRICS_PREFIX); + + S3StreamKafkaMetricsManager.configure(new MetricsConfig(metricsLevel, Attributes.empty(), metricsIntervalMs)); + S3StreamKafkaMetricsManager.initMetrics(meter, KAFKA_METRICS_PREFIX); + + TableTopicMetricsManager.initMetrics(meter); + } + + private static MetricsLevel parseMetricsLevel(String rawLevel) { + if (StringUtils.isBlank(rawLevel)) { + return MetricsLevel.INFO; + } + + try { + return MetricsLevel.valueOf(rawLevel.trim().toUpperCase(Locale.ENGLISH)); + } catch (IllegalArgumentException e) { + LOGGER.warn("Illegal metrics level '{}', defaulting to INFO", rawLevel); + return MetricsLevel.INFO; + } + } + + private static String buildJmxConfigPaths(KafkaConfig config) { + List paths = new ArrayList<>(); + paths.add(COMMON_JMX_PATH); + + Set roles = config.processRoles(); + if (roles.contains(ProcessRole.BrokerRole)) { + paths.add(BROKER_JMX_PATH); + } + if (roles.contains(ProcessRole.ControllerRole)) { + paths.add(CONTROLLER_JMX_PATH); + } + + return String.join(",", paths); + } +} diff --git a/core/src/main/java/org/apache/iceberg/avro/CodecSetup.java b/core/src/main/java/org/apache/iceberg/avro/CodecSetup.java new file mode 100644 index 0000000000..83ec87134c --- /dev/null +++ b/core/src/main/java/org/apache/iceberg/avro/CodecSetup.java @@ -0,0 +1,37 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.iceberg.avro; + +import org.apache.avro.LogicalTypes; + +public class CodecSetup { + + public static LogicalMap getLogicalMap() { + return LogicalMap.get(); + } + + static { + LogicalTypes.register(LogicalMap.NAME, schema -> LogicalMap.get()); + } + + public static void setup() { + } + +} diff --git a/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/ConsistentHashRouter.java b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/ConsistentHashRouter.java new file mode 100644 index 0000000000..f945595e59 --- /dev/null +++ b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/ConsistentHashRouter.java @@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package thirdparty.com.github.jaskey.consistenthash; + +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.Collection; +import java.util.Iterator; +import java.util.Map; +import java.util.SortedMap; +import java.util.TreeMap; + +/** + * To hash Node objects to a hash ring with a certain amount of virtual node. + * Method routeNode will return a Node instance which the object key should be allocated to according to consistent hash algorithm + * + * @param + */ +public class ConsistentHashRouter { + private final SortedMap> ring = new TreeMap<>(); + private final HashFunction hashFunction; + + public ConsistentHashRouter() { + this(null, 1); + } + + public ConsistentHashRouter(Collection pNodes, int vNodeCount) { + this(pNodes, vNodeCount, new MD5Hash()); + } + + /** + * + * @param pNodes collections of physical nodes + * @param vNodeCount amounts of virtual nodes + * @param hashFunction hash Function to hash Node instances + */ + public ConsistentHashRouter(Collection pNodes, int vNodeCount, HashFunction hashFunction) { + if (hashFunction == null) { + throw new NullPointerException("Hash Function is null"); + } + this.hashFunction = hashFunction; + if (pNodes != null) { + for (T pNode : pNodes) { + addNode(pNode, vNodeCount); + } + } + } + + /** + * add physic node to the hash ring with some virtual nodes + * @param pNode physical node needs added to hash ring + * @param vNodeCount the number of virtual node of the physical node. Value should be greater than or equals to 0 + */ + public void addNode(T pNode, int vNodeCount) { + if (vNodeCount < 0) throw new IllegalArgumentException("illegal virtual node counts :" + vNodeCount); + int existingReplicas = getExistingReplicas(pNode); + for (int i = 0; i < vNodeCount; i++) { + VirtualNode vNode = new VirtualNode<>(pNode, i + existingReplicas); + ring.put(hashFunction.hash(vNode.getKey()), vNode); + } + } + + /** + * remove the physical node from the hash ring + * @param pNode + */ + public void removeNode(T pNode) { + Iterator>> it = ring.entrySet().iterator(); + while (it.hasNext()) { + Map.Entry> entry = it.next(); + VirtualNode virtualNode = entry.getValue(); + if (virtualNode.isVirtualNodeOf(pNode)) { + it.remove(); + } + } + } + + /** + * with a specified key, route the nearest Node instance in the current hash ring + * @param objectKey the object key to find a nearest Node + * @return + */ + public T routeNode(String objectKey) { + if (ring.isEmpty()) { + return null; + } + Long hashVal = hashFunction.hash(objectKey); + SortedMap> tailMap = ring.tailMap(hashVal); + Long nodeHashVal = !tailMap.isEmpty() ? tailMap.firstKey() : ring.firstKey(); + return ring.get(nodeHashVal).getPhysicalNode(); + } + + + public int getExistingReplicas(T pNode) { + int replicas = 0; + for (VirtualNode vNode : ring.values()) { + if (vNode.isVirtualNodeOf(pNode)) { + replicas++; + } + } + return replicas; + } + + + //default hash function + private static class MD5Hash implements HashFunction { + MessageDigest instance; + + public MD5Hash() { + try { + instance = MessageDigest.getInstance("MD5"); + } catch (NoSuchAlgorithmException e) { + } + } + + @Override + public long hash(String key) { + instance.reset(); + instance.update(key.getBytes(StandardCharsets.UTF_8)); + byte[] digest = instance.digest(); + + long h = 0; + for (int i = 0; i < 4; i++) { + h <<= 8; + h |= ((int) digest[i]) & 0xFF; + } + return h; + } + } + +} \ No newline at end of file diff --git a/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/HashFunction.java b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/HashFunction.java new file mode 100644 index 0000000000..b60ac112bc --- /dev/null +++ b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/HashFunction.java @@ -0,0 +1,21 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package thirdparty.com.github.jaskey.consistenthash; + +public interface HashFunction { + long hash(String key); +} \ No newline at end of file diff --git a/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/Node.java b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/Node.java new file mode 100644 index 0000000000..43c91ef81f --- /dev/null +++ b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/Node.java @@ -0,0 +1,25 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package thirdparty.com.github.jaskey.consistenthash; + +public interface Node { + /** + * + * @return the key which will be used for hash mapping + */ + String getKey(); +} \ No newline at end of file diff --git a/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/VirtualNode.java b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/VirtualNode.java new file mode 100644 index 0000000000..21ac672a1f --- /dev/null +++ b/core/src/main/java/thirdparty/com/github/jaskey/consistenthash/VirtualNode.java @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package thirdparty.com.github.jaskey.consistenthash; + +public class VirtualNode implements Node { + final T physicalNode; + final int replicaIndex; + + public VirtualNode(T physicalNode, int replicaIndex) { + this.replicaIndex = replicaIndex; + this.physicalNode = physicalNode; + } + + @Override + public String getKey() { + return physicalNode.getKey() + "-" + replicaIndex; + } + + public boolean isVirtualNodeOf(T pNode) { + return physicalNode.getKey().equals(pNode.getKey()); + } + + public T getPhysicalNode() { + return physicalNode; + } +} \ No newline at end of file diff --git a/core/src/main/resources/jmx/rules/broker.yaml b/core/src/main/resources/jmx/rules/broker.yaml index ee2b066388..95576f5413 100644 --- a/core/src/main/resources/jmx/rules/broker.yaml +++ b/core/src/main/resources/jmx/rules/broker.yaml @@ -207,6 +207,13 @@ rules: type: gauge desc: Size of the request queue + - bean: kafka.network:type=RequestChannel,name=AvailableRequestSize + mapping: + Value: + metric: kafka.available.request.size + type: gauge + desc: Remaining permitted request size in the request queue + - bean: kafka.network:type=RequestChannel,name=ResponseQueueSize mapping: Value: @@ -370,4 +377,4 @@ rules: connection-accept-throttle-time: metric: kafka.listener.connection.accept.throttle.time type: gauge - desc: The average throttle-time pre listener \ No newline at end of file + desc: The average throttle-time pre listener diff --git a/core/src/main/scala/kafka/Kafka.scala b/core/src/main/scala/kafka/Kafka.scala index a5b14b11ea..865a0a78d5 100755 --- a/core/src/main/scala/kafka/Kafka.scala +++ b/core/src/main/scala/kafka/Kafka.scala @@ -17,8 +17,9 @@ package kafka +import com.automq.log.S3RollingFileAppender +import com.automq.opentelemetry.exporter.MetricsExportConfig import com.automq.shell.AutoMQApplication -import com.automq.shell.log.{LogUploader, S3LogConfig} import com.automq.stream.s3.ByteBufAlloc import joptsimple.OptionParser import kafka.autobalancer.metricsreporter.AutoBalancerMetricsReporter @@ -76,8 +77,7 @@ object Kafka extends Logging { private def enableApiForwarding(config: KafkaConfig) = config.migrationEnabled && config.interBrokerProtocolVersion.isApiForwardingEnabled - private def buildServer(props: Properties): Server = { - val config = KafkaConfig.fromProps(props, doLog = false) + private def buildServer(config: KafkaConfig): Server = { // AutoMQ for Kafka inject start // set allocator's policy as early as possible ByteBufAlloc.setPolicy(config.s3StreamAllocatorPolicy) @@ -89,18 +89,24 @@ object Kafka extends Logging { threadNamePrefix = None, enableForwarding = enableApiForwarding(config) ) + // AutoMQ for Kafka inject start AutoMQApplication.setClusterId(kafkaServer.clusterId) - AutoMQApplication.registerSingleton(classOf[S3LogConfig], new KafkaS3LogConfig(config, kafkaServer, null)) + S3RollingFileAppender.setup(new KafkaS3LogConfig(config, kafkaServer, null)) + AutoMQApplication.registerSingleton(classOf[MetricsExportConfig], new KafkaMetricsExportConfig(config, kafkaServer, null)) kafkaServer + // AutoMQ for Kafka inject end } else { val kafkaRaftServer = new KafkaRaftServer( config, Time.SYSTEM, ) + // AutoMQ for Kafka inject start AutoMQApplication.setClusterId(kafkaRaftServer.getSharedServer().clusterId) - AutoMQApplication.registerSingleton(classOf[S3LogConfig], new KafkaS3LogConfig(config, null, kafkaRaftServer)) + S3RollingFileAppender.setup(new KafkaS3LogConfig(config, null, kafkaRaftServer)) + AutoMQApplication.registerSingleton(classOf[MetricsExportConfig], new KafkaMetricsExportConfig(config, null, kafkaRaftServer)) AutoMQApplication.registerSingleton(classOf[KafkaRaftServer], kafkaRaftServer) kafkaRaftServer + // AutoMQ for Kafka inject end } } @@ -124,7 +130,8 @@ object Kafka extends Logging { val serverProps = getPropsFromArgs(args) addDefaultProps(serverProps) StorageUtil.formatStorage(serverProps) - val server = buildServer(serverProps) + val kafkaConfig = KafkaConfig.fromProps(serverProps, doLog = false) + val server = buildServer(kafkaConfig) AutoMQApplication.registerSingleton(classOf[Server], server) // AutoMQ for Kafka inject end @@ -141,7 +148,7 @@ object Kafka extends Logging { Exit.addShutdownHook("kafka-shutdown-hook", { try { server.shutdown() - LogUploader.getInstance().close() + S3RollingFileAppender.shutdown() } catch { case _: Throwable => fatal("Halting Kafka.") @@ -157,7 +164,6 @@ object Kafka extends Logging { fatal("Exiting Kafka due to fatal exception during startup.", e) Exit.exit(1) } - server.awaitShutdown() } catch { diff --git a/core/src/main/scala/kafka/KafkaMetricsExportConfig.scala b/core/src/main/scala/kafka/KafkaMetricsExportConfig.scala new file mode 100644 index 0000000000..1158b8a18c --- /dev/null +++ b/core/src/main/scala/kafka/KafkaMetricsExportConfig.scala @@ -0,0 +1,71 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka + +import com.automq.opentelemetry.exporter.MetricsExportConfig +import com.automq.stream.s3.operator.{ObjectStorage, ObjectStorageFactory} +import kafka.server.{KafkaConfig, KafkaRaftServer, KafkaServer} +import org.apache.commons.lang3.tuple.Pair + +import java.util + +class KafkaMetricsExportConfig( + config: KafkaConfig, + kafkaServer: KafkaServer, + kafkaRaftServer: KafkaRaftServer + ) extends MetricsExportConfig { + + private val _objectStorage = if (config.automq.opsBuckets().isEmpty) { + null + } else { + ObjectStorageFactory.instance().builder(config.automq.opsBuckets().get(0)).threadPrefix("s3-metrics").build() + } + + override def clusterId(): String = { + if (kafkaServer != null) { + kafkaServer.clusterId + } else { + kafkaRaftServer.getSharedServer().clusterId + } + } + + override def isLeader: Boolean = { + if (kafkaServer != null) { + // For broker mode, typically only one node should upload metrics + // You can implement your own leader selection logic here + false + } else { + // For KRaft mode, only active controller uploads metrics + kafkaRaftServer.controller.exists(controller => controller.controller != null && controller.controller.isActive) + } + } + + override def nodeId(): Int = config.nodeId + + override def objectStorage(): ObjectStorage = { + _objectStorage + } + + override def baseLabels(): util.List[Pair[String, String]] = { + config.automq.baseLabels() + } + + override def intervalMs(): Int = config.s3ExporterReportIntervalMs +} diff --git a/core/src/main/scala/kafka/KafkaS3LogConfig.scala b/core/src/main/scala/kafka/KafkaS3LogConfig.scala index 00560a53a2..894741432a 100644 --- a/core/src/main/scala/kafka/KafkaS3LogConfig.scala +++ b/core/src/main/scala/kafka/KafkaS3LogConfig.scala @@ -1,25 +1,33 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka -import com.automq.shell.log.S3LogConfig +import com.automq.log.uploader.S3LogConfig import com.automq.stream.s3.operator.{ObjectStorage, ObjectStorageFactory} import kafka.server.{KafkaConfig, KafkaRaftServer, KafkaServer} class KafkaS3LogConfig( - config: KafkaConfig, - kafkaServer: KafkaServer, - kafkaRaftServer: KafkaRaftServer -) extends S3LogConfig { + config: KafkaConfig, + kafkaServer: KafkaServer, + kafkaRaftServer: KafkaRaftServer + ) extends S3LogConfig { private val _objectStorage = if (config.automq.opsBuckets().isEmpty) { null @@ -29,15 +37,6 @@ class KafkaS3LogConfig( override def isEnabled: Boolean = config.s3OpsTelemetryEnabled - override def isActiveController: Boolean = { - - if (kafkaServer != null) { - false - } else { - kafkaRaftServer.controller.exists(controller => controller.controller != null && controller.controller.isActive) - } - } - override def clusterId(): String = { if (kafkaServer != null) { kafkaServer.clusterId @@ -52,4 +51,11 @@ class KafkaS3LogConfig( _objectStorage } + override def isLeader: Boolean = { + if (kafkaServer != null) { + false + } else { + kafkaRaftServer.controller.exists(controller => controller.controller != null && controller.controller.isActive) + } + } } diff --git a/core/src/main/scala/kafka/cluster/Partition.scala b/core/src/main/scala/kafka/cluster/Partition.scala index e8b03ecb86..d56f4de622 100755 --- a/core/src/main/scala/kafka/cluster/Partition.scala +++ b/core/src/main/scala/kafka/cluster/Partition.scala @@ -79,6 +79,21 @@ trait PartitionListener { * that the partition was deleted but only that this broker does not host a replica of it any more. */ def onDeleted(partition: TopicPartition): Unit = {} + + // AutoMQ inject start + /** + * Called when the partition leader epoch is changed. + */ + def onNewLeaderEpoch(oldEpoch: Long, newEpoch: Long): Unit = {} + + /** + * Called when there is a new record append which maybe not persisted yet. + * @param partition + * @param offset the next offset. + */ + def onNewAppend(partition: TopicPartition, offset: Long): Unit = {} + // AutoMQ inject end + } trait AlterPartitionListener { @@ -336,6 +351,10 @@ class Partition(val topicPartition: TopicPartition, // If ReplicaAlterLogDir command is in progress, this is future location of the log @volatile var futureLog: Option[UnifiedLog] = None + // AutoMQ inject start + @volatile var snapshotRead = false + // AutoMQ inject end + // Partition listeners private val listeners = new CopyOnWriteArrayList[PartitionListener]() @@ -347,6 +366,24 @@ class Partition(val topicPartition: TopicPartition, } } + // AutoMQ inject start + private val newLeaderEpochListener = new PartitionListener { + override def onNewLeaderEpoch(oldEpoch: Long, newEpoch: Long): Unit = { + listeners.forEach { listener => + listener.onNewLeaderEpoch(oldEpoch, newEpoch) + } + } + } + + private val newAppendListener = new PartitionListener { + override def onNewAppend(partition: TopicPartition, offset: Long): Unit = { + listeners.forEach { listener => + listener.onNewAppend(partition, offset) + } + } + } + // AutoMQ inject end + /* Epoch of the controller that last changed the leader. This needs to be initialized correctly upon broker startup. * One way of doing that is through the controller's start replica state change command. When a new broker starts up * the controller sends it a start replica command containing the leader for each partition that the broker hosts. @@ -364,7 +401,7 @@ class Partition(val topicPartition: TopicPartition, metricsGroup.newGauge("ReplicasCount", () => if (isLeader) assignmentState.replicationFactor else 0, tags) metricsGroup.newGauge("LastStableOffsetLag", () => log.map(_.lastStableOffsetLag).getOrElse(0), tags) - // AutoMQ for Kafka inject start + // AutoMQ inject start private val enableTraceLog = isTraceEnabled private var closed: Boolean = false /** @@ -373,7 +410,8 @@ class Partition(val topicPartition: TopicPartition, * Used to return fast when fetching messages with `fetchOffset` equals to `confirmOffset` in [[checkFetchOffsetAndMaybeGetInfo]] */ private var confirmOffset: Long = -1L - // AutoMQ for Kafka inject end + private val appendListeners = new CopyOnWriteArrayList[PartitionAppendListener]() + // AutoMQ inject end def hasLateTransaction(currentTimeMs: Long): Boolean = leaderLogIfLocal.exists(_.hasLateTransaction(currentTimeMs)) @@ -432,6 +470,25 @@ class Partition(val topicPartition: TopicPartition, listeners.remove(listener) } + // AutoMQ inject start + def addAppendListener(listener: PartitionAppendListener): Unit = { + appendListeners.add(listener) + } + + def removeAppendListener(listener: PartitionAppendListener): Unit = { + appendListeners.remove(listener) + } + + def notifyAppendListener(records: MemoryRecords): Unit = { + try { + appendListeners.forEach(_.onAppend(topicPartition, records)) + } catch { + case e: Exception => + error(s"Error while notifying append listeners for partition $topicPartition", e) + } + } + // AutoMQ inject end + /** * Create the future replica if 1) the current replica is not in the given log directory and 2) the future replica * does not exist. This method assumes that the current replica has already been created. @@ -511,7 +568,9 @@ class Partition(val topicPartition: TopicPartition, elasticUnifiedLog.confirmOffsetChangeListener = Some(() => handleLeaderConfirmOffsetMove()) // just update LEO to HW since we only have one replica val initialHighWatermark = log.logEndOffset - log.updateHighWatermark(log.logEndOffset) + // the high watermark is the same as the log end offset + log.updateHighWatermark(initialHighWatermark) + snapshotRead = elasticUnifiedLog.snapshotRead info(s"Log loaded for partition $topicPartition with initial high watermark $initialHighWatermark") case _ => updateHighWatermark(log) @@ -705,7 +764,7 @@ class Partition(val topicPartition: TopicPartition, // AutoMQ for Kafka inject start - private def checkClosed(): Unit = { + def checkClosed(): Unit = { if (closed) { throw new NotLeaderOrFollowerException("Leader %d for partition %s on broker %d is already closed" .format(localBrokerId, topicPartition, localBrokerId)) @@ -720,8 +779,12 @@ class Partition(val topicPartition: TopicPartition, */ def close(): Unit = { info("Closing partition") - logManager.removeFromCurrentLogs(topicPartition) - ElasticLogManager.removeLog(topicPartition) + log match { + case Some(unifiedLog: ElasticUnifiedLog) => + logManager.removeFromCurrentLogs(topicPartition, unifiedLog) + ElasticLogManager.removeLog(topicPartition, unifiedLog) + case _ => + } inWriteLock(leaderIsrUpdateLock) { closed = true } @@ -881,8 +944,12 @@ class Partition(val topicPartition: TopicPartition, // We update the leader epoch and the leader epoch start offset iff the // leader epoch changed. + val oldLeaderEpoch = leaderEpoch leaderEpoch = partitionState.leaderEpoch leaderEpochStartOffsetOpt = Some(leaderEpochStartOffset) + // AutoMQ inject start + newLeaderEpochListener.onNewLeaderEpoch(oldLeaderEpoch, partitionState.leaderEpoch) + // AutoMQ inject end } else { stateChangeLogger.info(s"Skipped the become-leader state change for $topicPartition with topic id $topicId " + s"and partition state $partitionState since it is already the leader with leader epoch $leaderEpoch. " + @@ -1487,6 +1554,13 @@ class Partition(val topicPartition: TopicPartition, // AutoMQ for Kafka inject end leaderLogIfLocal match { case Some(leaderLog) => + // AutoMQ inject start + if (snapshotRead) { + throw new NotLeaderOrFollowerException("Cannot write snapshot read partition %s on broker %d" + .format(topicPartition, localBrokerId)) + } + // AutoMQ inject end + val minIsr = effectiveMinIsr(leaderLog) val inSyncSize = partitionState.isr.size @@ -1499,6 +1573,11 @@ class Partition(val topicPartition: TopicPartition, val info = leaderLog.appendAsLeader(records, leaderEpoch = this.leaderEpoch, origin, interBrokerProtocolVersion, requestLocal, verificationGuard) + // AutoMQ inject start + notifyAppendListener(records) + newAppendListener.onNewAppend(topicPartition, leaderLog.logEndOffset) + // AutoMQ inject end + // we may need to increment high watermark since ISR could be down to 1 (info, maybeIncrementLeaderHW(leaderLog)) @@ -2294,4 +2373,50 @@ class Partition(val topicPartition: TopicPartition, partitionString.append("; LeaderRecoveryState: " + partitionState.leaderRecoveryState) partitionString.toString } + + // AutoMQ injection start + def withReadLock[T](fun: => T): T = { + inReadLock(leaderIsrUpdateLock)(fun) + } + + def withWriteLock[T](fun: => T): T = { + inWriteLock(leaderIsrUpdateLock)(fun) + } + + def addLogEventListener(listener: LogEventListener): Unit = { + log.get.asInstanceOf[ElasticUnifiedLog].getLocalLog().logSegmentManager.addLogEventListener(listener) + } + + def snapshot(): PartitionSnapshot = { + inReadLock(leaderIsrUpdateLock) { + val snapshot = PartitionSnapshot.builder() + snapshot.leaderEpoch(leaderEpoch) + val log = this.log.get.asInstanceOf[ElasticUnifiedLog] + log.snapshot(snapshot) + snapshot.build() + } + } + + def snapshot(snapshot: PartitionSnapshot): Unit = { + inWriteLock(leaderIsrUpdateLock) { + try { + if (enableTraceLog) { + trace(s"apply snapshot partition $topic-$partitionId $snapshot") + } + leaderEpoch = snapshot.leaderEpoch + val log = this.log.get.asInstanceOf[ElasticUnifiedLog] + log.snapshot(snapshot) + log.getLocalLog().appendAckThread.submit(() => { + // async it to avoid deadlock + tryCompleteDelayedRequests() + null + }) + } catch { + case e: Throwable => + error(s"apply snapshot fail, snapshot=$snapshot", e) + throw e + } + } + } + // AutoMQ injection end } diff --git a/core/src/main/scala/kafka/cluster/PartitionAppendListener.java b/core/src/main/scala/kafka/cluster/PartitionAppendListener.java new file mode 100644 index 0000000000..ad6ea551a4 --- /dev/null +++ b/core/src/main/scala/kafka/cluster/PartitionAppendListener.java @@ -0,0 +1,29 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.cluster; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.record.MemoryRecords; + +public interface PartitionAppendListener { + + void onAppend(TopicPartition topicPartition, MemoryRecords records); + +} diff --git a/core/src/main/scala/kafka/controller/streamaspect/client/Context.java b/core/src/main/scala/kafka/controller/streamaspect/client/Context.java index 0c6b9a2539..edc4d68898 100644 --- a/core/src/main/scala/kafka/controller/streamaspect/client/Context.java +++ b/core/src/main/scala/kafka/controller/streamaspect/client/Context.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.controller.streamaspect.client; diff --git a/core/src/main/scala/kafka/controller/streamaspect/client/StreamClientFactoryProxy.java b/core/src/main/scala/kafka/controller/streamaspect/client/StreamClientFactoryProxy.java index f0833d8a23..ffd26f72e1 100644 --- a/core/src/main/scala/kafka/controller/streamaspect/client/StreamClientFactoryProxy.java +++ b/core/src/main/scala/kafka/controller/streamaspect/client/StreamClientFactoryProxy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.controller.streamaspect.client; diff --git a/core/src/main/scala/kafka/controller/streamaspect/client/s3/StreamClientFactory.java b/core/src/main/scala/kafka/controller/streamaspect/client/s3/StreamClientFactory.java index 44596c9dc8..500e728753 100644 --- a/core/src/main/scala/kafka/controller/streamaspect/client/s3/StreamClientFactory.java +++ b/core/src/main/scala/kafka/controller/streamaspect/client/s3/StreamClientFactory.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.controller.streamaspect.client.s3; @@ -18,8 +26,12 @@ import org.apache.kafka.controller.stream.StreamClient; import com.automq.stream.s3.Config; +import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorageFactory; +import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_BACKGROUND; +import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_KEY; + public class StreamClientFactory { /** @@ -27,14 +39,14 @@ public class StreamClientFactory { */ public static StreamClient get(Context context) { Config streamConfig = ConfigUtils.to(context.kafkaConfig); + ObjectStorage objectStorage = ObjectStorageFactory.instance().builder() + .buckets(streamConfig.dataBuckets()) + .tagging(streamConfig.objectTagging()) + .extension(EXTENSION_TYPE_KEY, EXTENSION_TYPE_BACKGROUND) + .build(); return StreamClient.builder() .streamConfig(streamConfig) - .objectStorage( - ObjectStorageFactory.instance() - .builder(streamConfig.dataBuckets().get(0)) - .tagging(streamConfig.objectTagging()) - .build() - ) + .objectStorage(objectStorage) .build(); } } diff --git a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala index 58a37e3333..25d662c693 100644 --- a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala +++ b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala @@ -52,7 +52,7 @@ import scala.math.max * used by its callback. The delayed callback may acquire the group lock * since the delayed operation is completed only if the group lock can be acquired. */ -private[group] class GroupCoordinator( +class GroupCoordinator( val brokerId: Int, val groupConfig: GroupConfig, val offsetConfig: OffsetConfig, diff --git a/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala b/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala index a53d4485fb..a3fd9a8ab8 100644 --- a/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala +++ b/core/src/main/scala/kafka/coordinator/group/GroupCoordinatorAdapter.scala @@ -20,7 +20,7 @@ import kafka.common.OffsetAndMetadata import kafka.server.{KafkaConfig, ReplicaManager, RequestLocal} import kafka.utils.Implicits.MapExtensionMethods import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid} -import org.apache.kafka.common.message.{ConsumerGroupDescribeResponseData, ConsumerGroupHeartbeatRequestData, ConsumerGroupHeartbeatResponseData, DeleteGroupsResponseData, DescribeGroupsResponseData, HeartbeatRequestData, HeartbeatResponseData, JoinGroupRequestData, JoinGroupResponseData, LeaveGroupRequestData, LeaveGroupResponseData, ListGroupsRequestData, ListGroupsResponseData, OffsetCommitRequestData, OffsetCommitResponseData, OffsetDeleteRequestData, OffsetDeleteResponseData, OffsetFetchRequestData, OffsetFetchResponseData, ShareGroupDescribeResponseData, ShareGroupHeartbeatRequestData, ShareGroupHeartbeatResponseData, SyncGroupRequestData, SyncGroupResponseData, TxnOffsetCommitRequestData, TxnOffsetCommitResponseData} +import org.apache.kafka.common.message.{AutomqUpdateGroupRequestData, AutomqUpdateGroupResponseData, ConsumerGroupDescribeResponseData, ConsumerGroupHeartbeatRequestData, ConsumerGroupHeartbeatResponseData, DeleteGroupsResponseData, DescribeGroupsResponseData, HeartbeatRequestData, HeartbeatResponseData, JoinGroupRequestData, JoinGroupResponseData, LeaveGroupRequestData, LeaveGroupResponseData, ListGroupsRequestData, ListGroupsResponseData, OffsetCommitRequestData, OffsetCommitResponseData, OffsetDeleteRequestData, OffsetDeleteResponseData, OffsetFetchRequestData, OffsetFetchResponseData, ShareGroupDescribeResponseData, ShareGroupHeartbeatRequestData, ShareGroupHeartbeatResponseData, SyncGroupRequestData, SyncGroupResponseData, TxnOffsetCommitRequestData, TxnOffsetCommitResponseData} import org.apache.kafka.common.metrics.Metrics import org.apache.kafka.common.protocol.{ApiKeys, Errors} import org.apache.kafka.common.record.RecordBatch @@ -60,8 +60,8 @@ object GroupCoordinatorAdapter { * GroupCoordinatorAdapter is a thin wrapper around kafka.coordinator.group.GroupCoordinator * that exposes the new org.apache.kafka.coordinator.group.GroupCoordinator interface. */ -private[group] class GroupCoordinatorAdapter( - private val coordinator: GroupCoordinator, +class GroupCoordinatorAdapter( + val coordinator: GroupCoordinator, private val time: Time ) extends org.apache.kafka.coordinator.group.GroupCoordinator { @@ -645,4 +645,19 @@ private[group] class GroupCoordinatorAdapter( s"The old group coordinator does not support ${ApiKeys.SHARE_GROUP_DESCRIBE.name} API." )) } + + /** + * Update consumer groups + * + * @param context The coordinator request context. + * @param request The AutomqUpdateGroupRequestData data. + * @param bufferSupplier The buffer supplier tight to the request thread. + * @return A future yielding the response. + * The error code(s) of the response are set to indicate the error(s) occurred during the execution. + */ + override def updateGroup(context: RequestContext, + request: AutomqUpdateGroupRequestData, + bufferSupplier: BufferSupplier): CompletableFuture[AutomqUpdateGroupResponseData] = { + FutureUtils.failedFuture(new UnsupportedOperationException("Not implemented")) + } } diff --git a/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala b/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala index 59bdadb631..d06bb4f9cd 100644 --- a/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala +++ b/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala @@ -127,7 +127,9 @@ private[group] case object Empty extends GroupState { private object GroupMetadata extends Logging { + // AutoMQ for Kafka inject start private val CommitOffset: String = "CommitOffset" + // AutoMQ for Kafka inject end def loadGroup(groupId: String, initialState: GroupState, @@ -226,6 +228,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState var newMemberAdded: Boolean = false + // AutoMQ for Kafka inject start private val metricsGroup = new KafkaMetricsGroup(GroupMetadata.getClass) private def recreateOffsetMetric(tp: TopicPartition): Unit = { removeOffsetMetric(tp) @@ -247,7 +250,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState metricsGroup.removeMetric(GroupMetadata.CommitOffset, Map("group" -> groupId, "topic" -> tp.topic, "partition" -> tp.partition.toString).asJava) } - + // AutoMQ for Kafka inject end def inLock[T](fun: => T): T = CoreUtils.inLock(lock)(fun) def is(groupState: GroupState): Boolean = state == groupState @@ -452,8 +455,10 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState assertValidTransition(groupState) state = groupState currentStateTimestamp = Some(time.milliseconds()) + // AutoMQ for Kafka inject start if (groupState == Dead) offsets.foreach(offset => removeOffsetMetric(offset._1)) + // AutoMQ for Kafka inject end } def selectProtocol: String = { @@ -643,10 +648,12 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState def initializeOffsets(offsets: collection.Map[TopicPartition, CommitRecordMetadataAndOffset], pendingTxnOffsets: Map[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]]): Unit = { + // AutoMQ for Kafka inject start offsets.forKeyValue { (topicPartition, _) => if (!this.offsets.contains(topicPartition)) recreateOffsetMetric(topicPartition) } + // AutoMQ for Kafka inject end this.offsets ++= offsets this.pendingTransactionalOffsetCommits ++= pendingTxnOffsets } @@ -657,8 +664,10 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState if (offsetWithCommitRecordMetadata.appendedBatchOffset.isEmpty) throw new IllegalStateException("Cannot complete offset commit write without providing the metadata of the record " + "in the log.") + // AutoMQ for Kafka inject start if (!offsets.contains(topicPartition)) recreateOffsetMetric(topicPartition) + // AutoMQ for Kafka inject end if (!offsets.contains(topicPartition) || offsets(topicPartition).olderThan(offsetWithCommitRecordMetadata)) offsets.put(topicPartition, offsetWithCommitRecordMetadata) } @@ -746,6 +755,10 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState val currentOffsetOpt = offsets.get(topicPartition) if (currentOffsetOpt.forall(_.olderThan(commitRecordMetadataAndOffset))) { + // AutoMQ for Kafka inject start + if (!offsets.contains(topicPartition)) + recreateOffsetMetric(topicPartition) + // AutoMQ for Kafka inject end trace(s"TxnOffsetCommit for producer $producerId and group $groupId with offset $commitRecordMetadataAndOffset " + "committed and loaded into the cache.") offsets.put(topicPartition, commitRecordMetadataAndOffset) @@ -781,9 +794,11 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState pendingOffsets.remove(topicPartition) } val removedOffset = offsets.remove(topicPartition) + // AutoMQ for Kafka inject start if (removedOffset.isDefined) { removeOffsetMetric(topicPartition) } + // AutoMQ for Kafka inject end removedOffset.map(topicPartition -> _.offsetAndMetadata) }.toMap } @@ -893,4 +908,3 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState } } - diff --git a/core/src/main/scala/kafka/log/LogCleaner.scala b/core/src/main/scala/kafka/log/LogCleaner.scala index 1a1d33cdae..c1a6ce23bd 100644 --- a/core/src/main/scala/kafka/log/LogCleaner.scala +++ b/core/src/main/scala/kafka/log/LogCleaner.scala @@ -17,31 +17,31 @@ package kafka.log -import java.io.{File, IOException} -import java.nio._ -import java.util.Date -import java.util.concurrent.TimeUnit import kafka.common._ -import kafka.log.LogCleaner.{CleanerRecopyPercentMetricName, DeadThreadCountMetricName, MaxBufferUtilizationPercentMetricName, MaxCleanTimeMetricName, MaxCompactionDelayMetricsName} +import kafka.log.LogCleaner._ import kafka.log.streamaspect.ElasticLogSegment import kafka.server.{BrokerReconfigurable, KafkaConfig} import kafka.utils.{Logging, Pool} -import org.apache.kafka.common.{KafkaException, TopicPartition} import org.apache.kafka.common.config.ConfigException import org.apache.kafka.common.errors.{CorruptRecordException, KafkaStorageException} import org.apache.kafka.common.record.MemoryRecords.RecordFilter import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention import org.apache.kafka.common.record._ import org.apache.kafka.common.utils.{BufferSupplier, Time} +import org.apache.kafka.common.{KafkaException, TopicPartition} import org.apache.kafka.server.config.ServerConfigs import org.apache.kafka.server.metrics.KafkaMetricsGroup import org.apache.kafka.server.util.ShutdownableThread -import org.apache.kafka.storage.internals.log.{AbortedTxn, CleanerConfig, LastRecord, LogDirFailureChannel, LogSegment, LogSegmentOffsetOverflowException, OffsetMap, SkimpyOffsetMap, TransactionIndex} +import org.apache.kafka.storage.internals.log._ import org.apache.kafka.storage.internals.utils.Throttler -import scala.jdk.CollectionConverters._ +import java.io.{File, IOException} +import java.nio._ +import java.util.Date +import java.util.concurrent.TimeUnit import scala.collection.mutable.ListBuffer import scala.collection.{Iterable, Seq, Set, mutable} +import scala.jdk.CollectionConverters._ import scala.util.control.ControlThrowable /** @@ -977,6 +977,18 @@ private[log] class Cleaner(val id: Int, private[log] def groupSegmentsBySize(segments: Iterable[LogSegment], maxSize: Int, maxIndexSize: Int, firstUncleanableOffset: Long): List[Seq[LogSegment]] = { var grouped = List[List[LogSegment]]() var segs = segments.toList + + // AutoMQ inject start + def isOffsetRangeValid(group: List[LogSegment]) = { + val offsetRange = lastOffsetForFirstSegment(segs, firstUncleanableOffset) - group.last.baseOffset + // For ElasticLogSegment, use a stricter offset range check (`< Int.MaxValue`) to prevent a potential overflow + // issue as described in https://github.com/AutoMQ/automq/issues/2717. + // For other segment types, the original less-strict check (`<= Int.MaxValue`) is retained. + if (group.last.isInstanceOf[ElasticLogSegment]) offsetRange < Int.MaxValue + else offsetRange <= Int.MaxValue + } + // AutoMQ inject end + while (segs.nonEmpty) { var group = List(segs.head) var logSize = segs.head.size.toLong @@ -990,7 +1002,9 @@ private[log] class Cleaner(val id: Int, //if first segment size is 0, we don't need to do the index offset range check. //this will avoid empty log left every 2^31 message. (segs.head.size == 0 || - lastOffsetForFirstSegment(segs, firstUncleanableOffset) - group.last.baseOffset <= Int.MaxValue)) { + // AutoMQ inject start + isOffsetRangeValid(group))) { + // AutoMQ inject end group = segs.head :: group logSize += segs.head.size indexSize += offsetIndexSize(segs.head) @@ -1224,6 +1238,9 @@ private[log] class Cleaner(val id: Int, } for (batch <- fetchDataInfo.records.batches().asScala) { checkDone(topicPartition) + + writeBuffer.clear(); + val records = MemoryRecords.readableRecords(batch.asInstanceOf[DefaultRecordBatch].buffer()) throttler.maybeThrottle(records.sizeInBytes) val result = records.filterTo(topicPartition, logCleanerFilter, writeBuffer, maxLogMessageSize, decompressionBufferSupplier) @@ -1241,7 +1258,13 @@ private[log] class Cleaner(val id: Int, dest.append(result.maxOffset, result.maxTimestamp, result.shallowOffsetOfMaxTimestamp(), retained) throttler.maybeThrottle(outputBuffer.limit()) } + + // Grow the write buffer if needed, avoid always allocate a new buffer. + if (result.outputBuffer.capacity() > this.writeBuffer.capacity()) { + this.writeBuffer = ByteBuffer.allocate(result.outputBuffer.capacity()) + } } + restoreBuffers() } private def buildOffsetMapForSegmentV2(topicPartition: TopicPartition, diff --git a/core/src/main/scala/kafka/log/LogManager.scala b/core/src/main/scala/kafka/log/LogManager.scala index 862ecc3aa8..3514c7f8ff 100755 --- a/core/src/main/scala/kafka/log/LogManager.scala +++ b/core/src/main/scala/kafka/log/LogManager.scala @@ -17,7 +17,7 @@ package kafka.log -import kafka.log.streamaspect.{ElasticLogManager, ElasticUnifiedLog} +import kafka.log.streamaspect.{ElasticLogManager, ElasticUnifiedLog, OpenHint} import java.io._ import java.nio.file.{Files, NoSuchFileException} @@ -164,6 +164,10 @@ class LogManager(logDirs: Seq[File], Map("logDirectory" -> dir.getAbsolutePath).asJava) } + // AutoMQ inject start + private val snapshotReadLogs = new Pool[TopicPartition, UnifiedLog]() + // AutoMQ inject end + /** * Create and check validity of the given directories that are not in the given offline directories, specifically: *

    @@ -722,10 +726,14 @@ class LogManager(logDirs: Seq[File], val logs = logsInDir(localLogsByDir, dir) // update the last flush point - debug(s"Updating recovery points at $dir") + if (isDebugEnabled) { + debug(s"Updating recovery points at $dir") + } checkpointRecoveryOffsetsInDir(dir, logs) - debug(s"Updating log start offsets at $dir") + if (isDebugEnabled) { + debug(s"Updating log start offsets at $dir") + } checkpointLogStartOffsetsInDir(dir, logs) // mark that the shutdown was clean by creating marker file for log dirs that: @@ -735,7 +743,9 @@ class LogManager(logDirs: Seq[File], if (hadCleanShutdownFlags.getOrDefault(logDirAbsolutePath, false) || loadLogsCompletedFlags.getOrDefault(logDirAbsolutePath, false)) { val cleanShutdownFileHandler = new CleanShutdownFileHandler(dir.getPath) - debug(s"Writing clean shutdown marker at $dir with broker epoch=$brokerEpoch") + if (isDebugEnabled) { + debug(s"Writing clean shutdown marker at $dir with broker epoch=$brokerEpoch") + } CoreUtils.swallow(cleanShutdownFileHandler.write(brokerEpoch), this) } } @@ -958,8 +968,26 @@ class LogManager(logDirs: Seq[File], def getLog(topicPartition: TopicPartition, isFuture: Boolean = false): Option[UnifiedLog] = { if (isFuture) Option(futureLogs.get(topicPartition)) - else + else { + // AutoMQ inject start + val log = currentLogs.get(topicPartition) + if (log != null) { + Option(log) + } else { + Option(snapshotReadLogs.get(topicPartition)) + } + // AutoMQ inject end + } + } + + def getLogWithoutFallback(topicPartition: TopicPartition, isFuture: Boolean = false, isSnapshotRead: Boolean = false) = { + if (isFuture) { + Option(futureLogs.get(topicPartition)) + } else if (isSnapshotRead) { + Option(snapshotReadLogs.get(topicPartition)) + } else { Option(currentLogs.get(topicPartition)) + } } /** @@ -1050,7 +1078,7 @@ class LogManager(logDirs: Seq[File], // Only Partition#makeLeader will create a new log, the ReplicaManager#asyncApplyDelta will ensure the same partition // sequentially operate. So it's safe without lock // logCreationOrDeletionLock synchronized { - val log = getLog(topicPartition, isFuture).getOrElse { + val log = getLogWithoutFallback(topicPartition, isFuture, OpenHint.isSnapshotRead).getOrElse { // create the log if it has not already been created in another thread val now = time.milliseconds() @@ -1087,16 +1115,22 @@ class LogManager(logDirs: Seq[File], UnifiedLog.logDirName(topicPartition) } - val logDir = logDirs - .iterator // to prevent actually mapping the whole list, lazy map - .map(createLogDirectory(_, logDirName)) - .find(_.isSuccess) - .getOrElse(Failure(new KafkaStorageException("No log directories available. Tried " + logDirs.map(_.getAbsolutePath).mkString(", ")))) - .get // If Failure, will throw + val logDir = { + if (ElasticLogManager.enabled()) { + new File(logDirs.iterator.next().getAbsolutePath, logDirName) + } else { + logDirs + .iterator // to prevent actually mapping the whole list, lazy map + .map(createLogDirectory(_, logDirName)) + .find(_.isSuccess) + .getOrElse(Failure(new KafkaStorageException("No log directories available. Tried " + logDirs.map(_.getAbsolutePath).mkString(", ")))) + .get // If Failure, will throw + } + } val config = fetchLogConfig(topicPartition.topic) val log = if (ElasticLogManager.enabled()) { - ElasticLogManager.getOrCreateLog(logDir, config, scheduler, time, maxTransactionTimeoutMs, producerStateManagerConfig, brokerTopicStats, producerIdExpirationCheckIntervalMs, logDirFailureChannel, topicId, leaderEpoch) + ElasticLogManager.createLog(logDir, config, scheduler, time, maxTransactionTimeoutMs, producerStateManagerConfig, brokerTopicStats, producerIdExpirationCheckIntervalMs, logDirFailureChannel, topicId, leaderEpoch) } else { UnifiedLog( dir = logDir, @@ -1117,8 +1151,15 @@ class LogManager(logDirs: Seq[File], if (isFuture) futureLogs.put(topicPartition, log) - else - currentLogs.put(topicPartition, log) + else { + // AutoMQ inject start + if (OpenHint.isSnapshotRead) { + snapshotReadLogs.put(topicPartition, log) + } else { + currentLogs.put(topicPartition, log) + } + // AutoMQ inject end + } info(s"Created log for partition $topicPartition in $logDir with properties ${config.overriddenConfigsAsLoggableString} cost ${time.milliseconds() - now}ms") // Remove the preferred log dir since it has already been satisfied @@ -1363,8 +1404,12 @@ class LogManager(logDirs: Seq[File], } // AutoMQ for Kafka inject start - def removeFromCurrentLogs(topicPartition: TopicPartition): Unit = { - removeLogAndMetrics(currentLogs, topicPartition) + def removeFromCurrentLogs(topicPartition: TopicPartition, log: ElasticUnifiedLog): Unit = { + if (log.snapshotRead) { + removeLogAndMetrics(snapshotReadLogs, topicPartition) + } else { + removeLogAndMetrics(currentLogs, topicPartition) + } } // AutoMQ for Kafka inject end diff --git a/core/src/main/scala/kafka/log/UnifiedLog.scala b/core/src/main/scala/kafka/log/UnifiedLog.scala index 9976150273..c232edd88e 100644 --- a/core/src/main/scala/kafka/log/UnifiedLog.scala +++ b/core/src/main/scala/kafka/log/UnifiedLog.scala @@ -135,14 +135,14 @@ class UnifiedLog(@volatile var logStartOffset: Long, * that this could result in disagreement between replicas depending on when they began replicating the log. * In the worst case, the LSO could be seen by a consumer to go backwards. */ - @volatile private var firstUnstableOffsetMetadata: Option[LogOffsetMetadata] = None + @volatile protected var firstUnstableOffsetMetadata: Option[LogOffsetMetadata] = None /* Keep track of the current high watermark in order to ensure that segments containing offsets at or above it are * not eligible for deletion. This means that the active segment is only eligible for deletion if the high watermark * equals the log end offset (which may never happen for a partition under consistent load). This is needed to * prevent the log start offset (which is exposed in fetch responses) from getting ahead of the high watermark. */ - @volatile private var highWatermarkMetadata: LogOffsetMetadata = new LogOffsetMetadata(logStartOffset) + @volatile protected var highWatermarkMetadata: LogOffsetMetadata = new LogOffsetMetadata(logStartOffset) @volatile var partitionMetadataFile: Option[PartitionMetadataFile] = None @@ -455,7 +455,7 @@ class UnifiedLog(@volatile var logStartOffset: Long, private var metricNames: Map[String, java.util.Map[String, String]] = Map.empty newMetrics() - private[log] def newMetrics(): Unit = { + def newMetrics(): Unit = { val tags = (Map("topic" -> topicPartition.topic, "partition" -> topicPartition.partition.toString) ++ (if (isFuture) Map("is-future" -> "true") else Map.empty)).asJava metricsGroup.newGauge(LogMetricNames.NumLogSegments, () => numberOfSegments, tags) @@ -1983,8 +1983,6 @@ class UnifiedLog(@volatile var logStartOffset: Long, minOneMessage: Boolean): CompletableFuture[FetchDataInfo] = { CompletableFuture.completedFuture(read(startOffset, maxLength, isolation, minOneMessage)) } - - // AutoMQ inject end } diff --git a/core/src/main/scala/kafka/log/stream/s3/ConfigUtils.java b/core/src/main/scala/kafka/log/stream/s3/ConfigUtils.java index 87b9a2d601..f431697c27 100644 --- a/core/src/main/scala/kafka/log/stream/s3/ConfigUtils.java +++ b/core/src/main/scala/kafka/log/stream/s3/ConfigUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3; @@ -22,10 +30,12 @@ public static Config to(KafkaConfig s) { AutoMQConfig config = s.automq(); return new Config() .nodeId(s.nodeId()) + .nodeEpoch(config.nodeEpoch()) .dataBuckets(config.dataBuckets()) .walConfig(config.walConfig()) .walCacheSize(s.s3WALCacheSize()) .walUploadThreshold(s.s3WALUploadThreshold()) + .walUploadIntervalMs(s.s3WALUploadIntervalMs()) .streamSplitSize(s.s3StreamSplitSize()) .objectBlockSize(s.s3ObjectBlockSize()) .objectPartSize(s.s3ObjectPartSize()) diff --git a/core/src/main/scala/kafka/log/stream/s3/ControllerKVClient.java b/core/src/main/scala/kafka/log/stream/s3/ControllerKVClient.java index a619774cb6..0bf3416765 100644 --- a/core/src/main/scala/kafka/log/stream/s3/ControllerKVClient.java +++ b/core/src/main/scala/kafka/log/stream/s3/ControllerKVClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3; @@ -55,7 +63,9 @@ public ControllerKVClient(ControllerRequestSender requestSender) { @Override public CompletableFuture putKVIfAbsent(KeyValue keyValue) { - LOGGER.trace("[ControllerKVClient]: Put KV if absent: {}", keyValue); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Put KV if absent: {}", keyValue); + } PutKVRequest request = new PutKVRequest() .setKey(keyValue.key().get()) .setValue(keyValue.value().get().array()); @@ -83,7 +93,9 @@ public Builder toRequestBuilder() { Errors code = Errors.forCode(response.errorCode()); switch (code) { case NONE: - LOGGER.trace("[ControllerKVClient]: Put KV if absent: {}, result: {}", keyValue, response); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Put KV if absent: {}, result: {}", keyValue, response); + } return ResponseHandleResult.withSuccess(Value.of(response.value())); case KEY_EXIST: LOGGER.warn("[ControllerKVClient]: Failed to Put KV if absent: {}, code: {}, key already exist", keyValue, code); @@ -99,7 +111,9 @@ public Builder toRequestBuilder() { @Override public CompletableFuture putKV(KeyValue keyValue) { - LOGGER.trace("[ControllerKVClient]: Put KV: {}", keyValue); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Put KV: {}", keyValue); + } PutKVRequest request = new PutKVRequest() .setKey(keyValue.key().get()) .setValue(keyValue.value().get().array()) @@ -128,7 +142,9 @@ public Builder toRequestBuilder() { Errors code = Errors.forCode(response.errorCode()); switch (code) { case NONE: - LOGGER.trace("[ControllerKVClient]: Put KV: {}, result: {}", keyValue, response); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Put KV: {}, result: {}", keyValue, response); + } return ResponseHandleResult.withSuccess(Value.of(response.value())); default: LOGGER.error("[ControllerKVClient]: Failed to Put KV: {}, code: {}, retry later", keyValue, code); @@ -141,7 +157,9 @@ public Builder toRequestBuilder() { @Override public CompletableFuture getKV(Key key) { - LOGGER.trace("[ControllerKVClient]: Get KV: {}", key); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Get KV: {}", key); + } GetKVRequest request = new GetKVRequest() .setKey(key.get()); WrapRequest req = new BatchRequest() { @@ -169,7 +187,9 @@ public Builder toRequestBuilder() { switch (code) { case NONE: Value val = Value.of(response.value()); - LOGGER.trace("[ControllerKVClient]: Get KV: {}, result: {}", key, response); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Get KV: {}, result: {}", key, response); + } return ResponseHandleResult.withSuccess(val); default: LOGGER.error("[ControllerKVClient]: Failed to Get KV: {}, code: {}, retry later", key, code); @@ -182,7 +202,9 @@ public Builder toRequestBuilder() { @Override public CompletableFuture delKV(Key key) { - LOGGER.trace("[ControllerKVClient]: Delete KV: {}", key); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Delete KV: {}", key); + } DeleteKVRequest request = new DeleteKVRequest() .setKey(key.get()); WrapRequest req = new BatchRequest() { @@ -210,7 +232,9 @@ public Builder toRequestBuilder() { Errors code = Errors.forCode(response.errorCode()); switch (code) { case NONE: - LOGGER.trace("[ControllerKVClient]: Delete KV: {}, result: {}", key, response); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("[ControllerKVClient]: Delete KV: {}, result: {}", key, response); + } return ResponseHandleResult.withSuccess(Value.of(response.value())); case KEY_NOT_EXIST: LOGGER.info("[ControllerKVClient]: Delete KV: {}, result: KEY_NOT_EXIST", key); diff --git a/core/src/main/scala/kafka/log/stream/s3/DefaultS3Client.java b/core/src/main/scala/kafka/log/stream/s3/DefaultS3Client.java index 1433f151ab..bbb23cc466 100644 --- a/core/src/main/scala/kafka/log/stream/s3/DefaultS3Client.java +++ b/core/src/main/scala/kafka/log/stream/s3/DefaultS3Client.java @@ -1,25 +1,40 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3; +import kafka.autobalancer.metricsreporter.metric.Derivator; import kafka.log.stream.s3.metadata.StreamMetadataManager; import kafka.log.stream.s3.network.ControllerRequestSender; +import kafka.log.stream.s3.node.NodeManager; +import kafka.log.stream.s3.node.NodeManagerStub; +import kafka.log.stream.s3.node.NoopNodeManager; import kafka.log.stream.s3.objects.ControllerObjectManager; import kafka.log.stream.s3.streams.ControllerStreamManager; +import kafka.log.stream.s3.wal.ConfirmWal; +import kafka.log.stream.s3.wal.DefaultWalFactory; import kafka.server.BrokerServer; import org.apache.kafka.image.MetadataImage; import org.apache.kafka.server.common.automq.AutoMQVersion; +import com.automq.stream.Context; import com.automq.stream.api.Client; import com.automq.stream.api.KVClient; import com.automq.stream.api.StreamClient; @@ -39,37 +54,48 @@ import com.automq.stream.s3.failover.HaltStorageFailureHandler; import com.automq.stream.s3.failover.StorageFailureHandlerChain; import com.automq.stream.s3.index.LocalStreamRangeIndexCache; +import com.automq.stream.s3.metrics.S3StreamMetricsManager; +import com.automq.stream.s3.metrics.stats.NetworkStats; import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; import com.automq.stream.s3.network.NetworkBandwidthLimiter; import com.automq.stream.s3.objects.ObjectManager; -import com.automq.stream.s3.operator.BucketURI; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorageFactory; import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.s3.wal.DefaultWalHandle; +import com.automq.stream.s3.wal.WalFactory; +import com.automq.stream.s3.wal.WalHandle; import com.automq.stream.s3.wal.WriteAheadLog; -import com.automq.stream.s3.wal.impl.block.BlockWALService; -import com.automq.stream.s3.wal.impl.object.ObjectWALConfig; -import com.automq.stream.s3.wal.impl.object.ObjectWALService; -import com.automq.stream.utils.IdURI; import com.automq.stream.utils.LogContext; -import com.automq.stream.utils.Time; import com.automq.stream.utils.threads.S3StreamThreadPoolMonitor; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collections; +import java.util.HashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; +import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_BACKGROUND; +import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_KEY; +import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_MAIN; + public class DefaultS3Client implements Client { private static final Logger LOGGER = LoggerFactory.getLogger(DefaultS3Client.class); protected final Config config; + protected final Derivator networkInboundRate = new Derivator(); + protected final Derivator networkOutboundRate = new Derivator(); private StreamMetadataManager metadataManager; protected ControllerRequestSender requestSender; + protected ObjectStorage mainObjectStorage; + protected ObjectStorage backgroundObjectStorage; + protected WriteAheadLog writeAheadLog; + protected StorageFailureHandlerChain storageFailureHandlerChain; protected S3Storage storage; protected ObjectReaderFactory objectReaderFactory; @@ -79,6 +105,8 @@ public class DefaultS3Client implements Client { protected StreamManager streamManager; + protected NodeManager nodeManager; + protected CompactionManager compactionManager; protected S3StreamClient streamClient; @@ -100,51 +128,53 @@ public DefaultS3Client(BrokerServer brokerServer, Config config) { @Override public void start() { - BucketURI dataBucket = config.dataBuckets().get(0); long refillToken = (long) (config.networkBaselineBandwidth() * ((double) config.refillPeriodMs() / 1000)); if (refillToken <= 0) { throw new IllegalArgumentException(String.format("refillToken must be greater than 0, bandwidth: %d, refill period: %dms", config.networkBaselineBandwidth(), config.refillPeriodMs())); } - long maxToken = refillToken * 10; GlobalNetworkBandwidthLimiters.instance().setup(AsyncNetworkBandwidthLimiter.Type.INBOUND, - refillToken, config.refillPeriodMs(), maxToken); + refillToken, config.refillPeriodMs(), config.networkBaselineBandwidth()); networkInboundLimiter = GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.INBOUND); + S3StreamMetricsManager.registerNetworkAvailableBandwidthSupplier(AsyncNetworkBandwidthLimiter.Type.INBOUND, () -> + config.networkBaselineBandwidth() - (long) networkInboundRate.derive( + TimeUnit.NANOSECONDS.toSeconds(System.nanoTime()), NetworkStats.getInstance().networkInboundUsageTotal().get())); + // Use a larger token pool for outbound traffic to avoid spikes caused by Upload WAL affecting tail-reading performance. GlobalNetworkBandwidthLimiters.instance().setup(AsyncNetworkBandwidthLimiter.Type.OUTBOUND, - refillToken, config.refillPeriodMs(), maxToken); + refillToken, config.refillPeriodMs(), config.networkBaselineBandwidth() * 5); networkOutboundLimiter = GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.OUTBOUND); - ObjectStorage objectStorage = ObjectStorageFactory.instance().builder(dataBucket).tagging(config.objectTagging()) - .inboundLimiter(networkInboundLimiter).outboundLimiter(networkOutboundLimiter).readWriteIsolate(true) - .threadPrefix("dataflow").build(); - if (!objectStorage.readinessCheck()) { + S3StreamMetricsManager.registerNetworkAvailableBandwidthSupplier(AsyncNetworkBandwidthLimiter.Type.OUTBOUND, () -> + config.networkBaselineBandwidth() - (long) networkOutboundRate.derive( + TimeUnit.NANOSECONDS.toSeconds(System.nanoTime()), NetworkStats.getInstance().networkOutboundUsageTotal().get())); + + this.localIndexCache = LocalStreamRangeIndexCache.create(); + this.objectReaderFactory = new DefaultObjectReaderFactory(() -> this.mainObjectStorage); + this.metadataManager = new StreamMetadataManager(brokerServer, config.nodeId(), objectReaderFactory, localIndexCache); + this.requestSender = new ControllerRequestSender(brokerServer, new ControllerRequestSender.RetryPolicyContext(config.controllerRequestRetryMaxCount(), + config.controllerRequestRetryBaseDelayMs())); + this.streamManager = newStreamManager(config.nodeId(), config.nodeEpoch(), false); + this.objectManager = newObjectManager(config.nodeId(), config.nodeEpoch(), false); + this.mainObjectStorage = newMainObjectStorage(); + if (!mainObjectStorage.readinessCheck()) { throw new IllegalArgumentException(String.format("%s is not ready", config.dataBuckets())); } - ObjectStorage compactionobjectStorage = ObjectStorageFactory.instance().builder(dataBucket).tagging(config.objectTagging()) - .inboundLimiter(networkInboundLimiter).outboundLimiter(networkOutboundLimiter) - .threadPrefix("compaction").build(); - ControllerRequestSender.RetryPolicyContext retryPolicyContext = new ControllerRequestSender.RetryPolicyContext(config.controllerRequestRetryMaxCount(), - config.controllerRequestRetryBaseDelayMs()); - localIndexCache = new LocalStreamRangeIndexCache(); - localIndexCache.init(config.nodeId(), objectStorage); + this.backgroundObjectStorage = newBackgroundObjectStorage(); + localIndexCache.init(config.nodeId(), backgroundObjectStorage); localIndexCache.start(); - this.objectReaderFactory = new DefaultObjectReaderFactory(objectStorage); - this.metadataManager = new StreamMetadataManager(brokerServer, config.nodeId(), objectReaderFactory, localIndexCache); - this.requestSender = new ControllerRequestSender(brokerServer, retryPolicyContext); - this.streamManager = newStreamManager(config.nodeId(), config.nodeEpoch(), false); this.streamManager.setStreamCloseHook(streamId -> localIndexCache.uploadOnStreamClose()); - this.objectManager = newObjectManager(config.nodeId(), config.nodeEpoch(), false); this.objectManager.setCommitStreamSetObjectHook(localIndexCache::updateIndexFromRequest); - this.blockCache = new StreamReaders(this.config.blockCacheSize(), objectManager, objectStorage, objectReaderFactory); - this.compactionManager = new CompactionManager(this.config, this.objectManager, this.streamManager, compactionobjectStorage); + this.blockCache = new StreamReaders(this.config.blockCacheSize(), objectManager, mainObjectStorage, objectReaderFactory); + this.compactionManager = new CompactionManager(this.config, this.objectManager, this.streamManager, backgroundObjectStorage); this.writeAheadLog = buildWAL(); - StorageFailureHandlerChain storageFailureHandler = new StorageFailureHandlerChain(); - this.storage = new S3Storage(this.config, writeAheadLog, streamManager, objectManager, blockCache, objectStorage, storageFailureHandler); + this.storageFailureHandlerChain = new StorageFailureHandlerChain(); + this.storage = newS3Storage(); // stream object compactions share the same object storage with stream set object compactions - this.streamClient = new S3StreamClient(this.streamManager, this.storage, this.objectManager, compactionobjectStorage, this.config, networkInboundLimiter, networkOutboundLimiter); - storageFailureHandler.addHandler(new ForceCloseStorageFailureHandler(streamClient)); - storageFailureHandler.addHandler(new HaltStorageFailureHandler()); + this.streamClient = new S3StreamClient(this.streamManager, this.storage, this.objectManager, backgroundObjectStorage, this.config, networkInboundLimiter, networkOutboundLimiter); + storageFailureHandlerChain.addHandler(new ForceCloseStorageFailureHandler(streamClient)); + storageFailureHandlerChain.addHandler(new HaltStorageFailureHandler()); this.streamClient.registerStreamLifeCycleListener(localIndexCache); this.kvClient = new ControllerKVClient(this.requestSender); + Context.instance().kvClient(this.kvClient); this.failover = failover(); S3StreamThreadPoolMonitor.config(new LogContext("ThreadPoolMonitor").logger("s3.threads.logger"), TimeUnit.SECONDS.toMillis(5)); @@ -182,29 +212,34 @@ public CompletableFuture failover(FailoverRequest request) { } protected WriteAheadLog buildWAL() { - IdURI uri = IdURI.parse(config.walConfig()); - switch (uri.protocol()) { - case "file": - return BlockWALService.builder(uri).config(config).build(); - case "s3": - ObjectStorage walObjectStorage = ObjectStorageFactory.instance() - .builder(BucketURI.parse(config.walConfig())) - .inboundLimiter(networkInboundLimiter) - .outboundLimiter(networkOutboundLimiter) - .tagging(config.objectTagging()) - .threadPrefix("s3-wal") - .build(); - - ObjectWALConfig.Builder configBuilder = ObjectWALConfig.builder() - .withURI(uri) - .withClusterId(brokerServer.clusterId()) - .withNodeId(config.nodeId()) - .withEpoch(config.nodeEpoch()); - - return new ObjectWALService(Time.SYSTEM, walObjectStorage, configBuilder.build()); - default: - throw new IllegalArgumentException("Invalid WAL schema: " + uri.protocol()); - } + String clusterId = brokerServer.clusterId(); + WalHandle walHandle = new DefaultWalHandle(clusterId); + WalFactory factory = new DefaultWalFactory(config.nodeId(), config.objectTagging(), networkInboundLimiter, networkOutboundLimiter); + return new ConfirmWal(config.nodeId(), config.nodeEpoch(), config.walConfig(), false, factory, getNodeManager(), walHandle); + } + + protected ObjectStorage newMainObjectStorage() { + return ObjectStorageFactory.instance().builder() + .buckets(config.dataBuckets()) + .tagging(config.objectTagging()) + .extension(EXTENSION_TYPE_KEY, EXTENSION_TYPE_MAIN) + .readWriteIsolate(true) + .inboundLimiter(networkInboundLimiter) + .outboundLimiter(networkOutboundLimiter) + .threadPrefix("main") + .build(); + } + + protected ObjectStorage newBackgroundObjectStorage() { + return ObjectStorageFactory.instance().builder() + .buckets(config.dataBuckets()) + .tagging(config.objectTagging()) + .extension(EXTENSION_TYPE_KEY, EXTENSION_TYPE_BACKGROUND) + .readWriteIsolate(false) + .inboundLimiter(networkInboundLimiter) + .outboundLimiter(networkOutboundLimiter) + .threadPrefix("background") + .build(); } protected StreamManager newStreamManager(int nodeId, long nodeEpoch, boolean failoverMode) { @@ -217,6 +252,10 @@ protected ObjectManager newObjectManager(int nodeId, long nodeEpoch, boolean fai this::getAutoMQVersion, failoverMode); } + protected S3Storage newS3Storage() { + return new S3Storage(config, writeAheadLog, streamManager, objectManager, blockCache, mainObjectStorage, storageFailureHandlerChain); + } + protected Failover failover() { return new Failover(new FailoverFactory() { @Override @@ -231,7 +270,13 @@ public ObjectManager getObjectManager(int nodeId, long nodeEpoch) { @Override public WriteAheadLog getWal(FailoverRequest request) { - return BlockWALService.recoveryBuilder(request.getDevice()).build(); + String clusterId = brokerServer.clusterId(); + int nodeId = request.getNodeId(); + long nodeEpoch = request.getNodeEpoch(); + WalHandle walHandle = new DefaultWalHandle(clusterId); + WalFactory factory = new DefaultWalFactory(nodeId, config.objectTagging(), networkInboundLimiter, networkOutboundLimiter); + NodeManager nodeManager = new NodeManagerStub(requestSender, nodeId, nodeEpoch, Collections.emptyMap()); + return new ConfirmWal(nodeId, nodeEpoch, request.getKraftWalConfigs(), true, factory, nodeManager, walHandle); } }, (wal, sm, om, logger) -> { try { @@ -248,4 +293,13 @@ protected AutoMQVersion getAutoMQVersion() { } return brokerServer.metadataCache().autoMQVersion(); } + + private NodeManager getNodeManager() { + if (this.nodeManager == null) { + this.nodeManager = config.version().isWalRegistrationSupported() + ? new NodeManagerStub(this.requestSender, config.nodeId(), config.nodeEpoch(), new HashMap<>()) + : new NoopNodeManager(config.nodeId(), config.nodeEpoch()); + } + return this.nodeManager; + } } diff --git a/core/src/main/scala/kafka/log/stream/s3/metadata/InRangeObjectsFetcher.java b/core/src/main/scala/kafka/log/stream/s3/metadata/InRangeObjectsFetcher.java index 0b32501402..78a7ce05ee 100644 --- a/core/src/main/scala/kafka/log/stream/s3/metadata/InRangeObjectsFetcher.java +++ b/core/src/main/scala/kafka/log/stream/s3/metadata/InRangeObjectsFetcher.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.metadata; diff --git a/core/src/main/scala/kafka/log/stream/s3/metadata/StreamMetadataManager.java b/core/src/main/scala/kafka/log/stream/s3/metadata/StreamMetadataManager.java index 4676a13faf..59bfee341e 100644 --- a/core/src/main/scala/kafka/log/stream/s3/metadata/StreamMetadataManager.java +++ b/core/src/main/scala/kafka/log/stream/s3/metadata/StreamMetadataManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.metadata; @@ -37,7 +45,9 @@ import com.automq.stream.s3.objects.ObjectAttributes; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorage.ReadOptions; +import com.automq.stream.s3.streams.StreamMetadataListener; import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Threads; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -45,11 +55,12 @@ import java.util.ArrayList; import java.util.LinkedList; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.stream.Collectors; import io.netty.buffer.ByteBuf; @@ -65,14 +76,17 @@ public class StreamMetadataManager implements InRangeObjectsFetcher, MetadataPub private MetadataImage metadataImage; private final ObjectReaderFactory objectReaderFactory; private final LocalStreamRangeIndexCache indexCache; + private final Map streamMetadataListeners = new ConcurrentHashMap<>(); - public StreamMetadataManager(BrokerServer broker, int nodeId, ObjectReaderFactory objectReaderFactory, LocalStreamRangeIndexCache indexCache) { + public StreamMetadataManager(BrokerServer broker, int nodeId, ObjectReaderFactory objectReaderFactory, + LocalStreamRangeIndexCache indexCache) { this.nodeId = nodeId; - this.metadataImage = broker.metadataCache().currentImage(); + this.metadataImage = broker.metadataCache().retainedImage(); this.pendingGetObjectsTasks = new LinkedList<>(); this.objectReaderFactory = objectReaderFactory; this.indexCache = indexCache; - this.pendingExecutorService = Executors.newSingleThreadScheduledExecutor(new DefaultThreadFactory("pending-get-objects-task-executor")); + this.pendingExecutorService = + Threads.newSingleThreadScheduledExecutor(new DefaultThreadFactory("pending-get-objects-task-executor"), LOGGER); broker.metadataLoader().installPublishers(List.of(this)).join(); } @@ -83,15 +97,21 @@ public String name() { @Override public void onMetadataUpdate(MetadataDelta delta, MetadataImage newImage, LoaderManifest manifest) { + Set changedStreams; synchronized (this) { if (newImage.highestOffsetAndEpoch().equals(this.metadataImage.highestOffsetAndEpoch())) { return; } + newImage.retain(); + MetadataImage oldImage = this.metadataImage; this.metadataImage = newImage; + changedStreams = delta.getOrCreateStreamsMetadataDelta().changedStreams(); + oldImage.release(); } // retry all pending tasks retryPendingTasks(); this.indexCache.asyncPrune(this::getStreamSetObjectIds); + notifyMetadataListeners(changedStreams); } public CompletableFuture> getStreamSetObjects() { @@ -122,12 +142,12 @@ public Set getStreamSetObjectIds() { public CompletableFuture fetch(long streamId, long startOffset, long endOffset, int limit) { // TODO: cache the object list for next search CompletableFuture cf = new CompletableFuture<>(); - exec(() -> fetch0(cf, streamId, startOffset, endOffset, limit), cf, LOGGER, "fetchObjects"); + exec(() -> fetch0(cf, streamId, startOffset, endOffset, limit, false), cf, LOGGER, "fetchObjects"); return cf; } private void fetch0(CompletableFuture cf, long streamId, - long startOffset, long endOffset, int limit) { + long startOffset, long endOffset, int limit, boolean retryFetch) { Image image = getImage(); try { final S3StreamsMetadataImage streamsImage = image.streamsMetadata(); @@ -163,9 +183,11 @@ private void fetch0(CompletableFuture cf, long streamId, streamId, startOffset, endOffset, limit, rst.objects().size(), rst.endOffset()); CompletableFuture pendingCf = pendingFetch(); - pendingCf.thenAccept(nil -> fetch0(cf, streamId, startOffset, endOffset, limit)); - cf.whenComplete((r, ex) -> - LOGGER.info("[FetchObjects],[COMPLETE_PENDING],streamId={} startOffset={} endOffset={} limit={}", streamId, startOffset, endOffset, limit)); + pendingCf.thenAccept(nil -> fetch0(cf, streamId, startOffset, endOffset, limit, true)); + if (!retryFetch) { + cf.whenComplete((r, ex) -> + LOGGER.info("[FetchObjects],[COMPLETE_PENDING],streamId={} startOffset={} endOffset={} limit={}", streamId, startOffset, endOffset, limit)); + } }).exceptionally(ex -> { cf.completeExceptionally(ex); return null; @@ -204,22 +226,22 @@ public List getStreamMetadataList(List streamIds) { try (Image image = getImage()) { final S3StreamsMetadataImage streamsImage = image.streamsMetadata(); - List streamMetadataList = new ArrayList<>(); - for (Long streamId : streamIds) { - S3StreamMetadataImage streamImage = streamsImage.timelineStreamMetadata().get(streamId); - if (streamImage == null) { - LOGGER.warn("[GetStreamMetadataList]: stream: {} not exists", streamId); - continue; - } - StreamMetadata streamMetadata = new StreamMetadata(streamId, streamImage.getEpoch(), - streamImage.getStartOffset(), -1L, streamImage.state()) { - @Override - public long endOffset() { - throw new UnsupportedOperationException(); + List streamMetadataList = new ArrayList<>(streamIds.size()); + streamsImage.inLockRun(() -> { + for (Long streamId : streamIds) { + S3StreamMetadataImage streamImage = streamsImage.timelineStreamMetadata().get(streamId); + if (streamImage == null) { + LOGGER.warn("[GetStreamMetadataList]: stream: {} not exists", streamId); + continue; } - }; - streamMetadataList.add(streamMetadata); - } + // If there is a streamImage, it means the stream exists. + @SuppressWarnings("OptionalGetWithoutIsPresent") long endOffset = streamsImage.streamEndOffset(streamId).getAsLong(); + StreamMetadata streamMetadata = new StreamMetadata(streamId, streamImage.getEpoch(), + streamImage.getStartOffset(), endOffset, streamImage.state()); + Optional.ofNullable(streamImage.lastRange()).ifPresent(r -> streamMetadata.nodeId(r.nodeId())); + streamMetadataList.add(streamMetadata); + } + }); return streamMetadataList; } } @@ -242,6 +264,27 @@ public int getObjectsCount() { } } + public synchronized StreamMetadataListener.Handle addMetadataListener(long streamId, StreamMetadataListener listener) { + streamMetadataListeners.put(streamId, listener); + List list = getStreamMetadataList(List.of(streamId)); + if (!list.isEmpty()) { + listener.onNewStreamMetadata(list.get(0)); + } + return () -> streamMetadataListeners.remove(streamId, listener); + } + + private synchronized void notifyMetadataListeners(Set changedStreams) { + changedStreams.forEach(streamId -> { + StreamMetadataListener listener = streamMetadataListeners.get(streamId); + if (listener != null) { + List list = getStreamMetadataList(List.of(streamId)); + if (!list.isEmpty()) { + listener.onNewStreamMetadata(list.get(0)); + } + } + }); + } + // must access thread safe private CompletableFuture pendingFetch() { GetObjectsTask task = new GetObjectsTask(); diff --git a/core/src/main/scala/kafka/log/stream/s3/network/ControllerRequestSender.java b/core/src/main/scala/kafka/log/stream/s3/network/ControllerRequestSender.java index d3dc578628..3eb6c6a8c3 100644 --- a/core/src/main/scala/kafka/log/stream/s3/network/ControllerRequestSender.java +++ b/core/src/main/scala/kafka/log/stream/s3/network/ControllerRequestSender.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.network; @@ -23,6 +31,9 @@ import org.apache.kafka.server.ControllerRequestCompletionHandler; import org.apache.kafka.server.NodeToControllerChannelManager; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Threads; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -31,18 +42,18 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; -import io.netty.util.concurrent.DefaultThreadFactory; - public class ControllerRequestSender { + private static final Logger LOGGER = LoggerFactory.getLogger(ControllerRequestSender.class); - private static final long MAX_RETRY_DELAY_MS = 10 * 1000; // 10s + + private static final long MAX_RETRY_DELAY_MS = Systems.getEnvLong("AUTOMQ_CONTROLLER_REQUEST_MAX_RETRY_DELAY_MS", 10L * 1000); // 10s + private final RetryPolicyContext retryPolicyContext; private final NodeToControllerChannelManager channelManager; @@ -55,7 +66,8 @@ public ControllerRequestSender(BrokerServer brokerServer, RetryPolicyContext ret this.retryPolicyContext = retryPolicyContext; this.channelManager = brokerServer.newNodeToControllerChannelManager("s3stream-to-controller", 60000); this.channelManager.start(); - this.retryService = Executors.newSingleThreadScheduledExecutor(new DefaultThreadFactory("controller-request-retry-sender")); + this.retryService = + Threads.newSingleThreadScheduledExecutor("controller-request-retry-sender", false, LOGGER); this.requestAccumulatorMap = new ConcurrentHashMap<>(); } diff --git a/core/src/main/scala/kafka/log/stream/s3/network/request/BatchRequest.java b/core/src/main/scala/kafka/log/stream/s3/network/request/BatchRequest.java index 2887072b36..bdd96cce7e 100644 --- a/core/src/main/scala/kafka/log/stream/s3/network/request/BatchRequest.java +++ b/core/src/main/scala/kafka/log/stream/s3/network/request/BatchRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.network.request; diff --git a/core/src/main/scala/kafka/log/stream/s3/network/request/WrapRequest.java b/core/src/main/scala/kafka/log/stream/s3/network/request/WrapRequest.java index 8b46e3a376..af6d0c1ef5 100644 --- a/core/src/main/scala/kafka/log/stream/s3/network/request/WrapRequest.java +++ b/core/src/main/scala/kafka/log/stream/s3/network/request/WrapRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.network.request; diff --git a/core/src/main/scala/kafka/log/stream/s3/node/NodeManager.java b/core/src/main/scala/kafka/log/stream/s3/node/NodeManager.java new file mode 100644 index 0000000000..ba1e688f17 --- /dev/null +++ b/core/src/main/scala/kafka/log/stream/s3/node/NodeManager.java @@ -0,0 +1,38 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.stream.s3.node; + +import org.apache.kafka.controller.stream.NodeMetadata; + +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +public interface NodeManager { + + CompletableFuture update(Function> updater); + + default CompletableFuture updateWal(String walConfig) { + return update(nodeMetadata -> Optional.of(new NodeMetadata(nodeMetadata.getNodeId(), nodeMetadata.getNodeEpoch(), walConfig, nodeMetadata.getTags()))); + } + + CompletableFuture getNodeMetadata(); + +} diff --git a/core/src/main/scala/kafka/log/stream/s3/node/NodeManagerStub.java b/core/src/main/scala/kafka/log/stream/s3/node/NodeManagerStub.java new file mode 100644 index 0000000000..0c8a45abb2 --- /dev/null +++ b/core/src/main/scala/kafka/log/stream/s3/node/NodeManagerStub.java @@ -0,0 +1,196 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.stream.s3.node; + +import kafka.log.stream.s3.network.ControllerRequestSender; +import kafka.log.stream.s3.network.request.WrapRequest; + +import org.apache.kafka.common.message.AutomqGetNodesRequestData; +import org.apache.kafka.common.message.AutomqGetNodesResponseData; +import org.apache.kafka.common.message.AutomqRegisterNodeRequestData; +import org.apache.kafka.common.message.AutomqRegisterNodeResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.AbstractRequest; +import org.apache.kafka.common.requests.s3.AutomqGetNodesRequest; +import org.apache.kafka.common.requests.s3.AutomqGetNodesResponse; +import org.apache.kafka.common.requests.s3.AutomqRegisterNodeRequest; +import org.apache.kafka.common.requests.s3.AutomqRegisterNodeResponse; +import org.apache.kafka.controller.stream.NodeMetadata; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static kafka.log.stream.s3.network.ControllerRequestSender.ResponseHandleResult.withRetry; +import static kafka.log.stream.s3.network.ControllerRequestSender.ResponseHandleResult.withSuccess; + +public class NodeManagerStub implements NodeManager { + private static final Logger LOGGER = LoggerFactory.getLogger(NodeManagerStub.class); + private final ControllerRequestSender requestSender; + private final int nodeId; + private final long epoch; + private NodeMetadata nodeMetadata; + private CompletableFuture lastCf; + + public NodeManagerStub(ControllerRequestSender requestSender, int nodeId, long epoch, + Map staticTags) { + this.requestSender = requestSender; + this.nodeId = nodeId; + this.epoch = epoch; + this.lastCf = getNodeMetadata0().thenAccept(opt -> { + synchronized (NodeManagerStub.this) { + nodeMetadata = opt.map(o -> { + Map newTags = new HashMap<>(o.getTags()); + newTags.putAll(staticTags); + return new NodeMetadata(o.getNodeId(), o.getNodeEpoch(), o.getWalConfig(), newTags); + }).orElseGet(() -> new NodeMetadata(nodeId, epoch, "", new HashMap<>(staticTags))); + } + }); + + } + + public synchronized CompletableFuture update(Function> updater) { + CompletableFuture cf = new CompletableFuture<>(); + lastCf.whenComplete((rst, e) -> { + if (e == null) { + update0(updater, cf); + } else { + cf.completeExceptionally(e); + } + }); + this.lastCf = cf; + return cf; + } + + public void update0(Function> updater, CompletableFuture cf) { + try { + Optional newNodeMetadataOpt = updater.apply(nodeMetadata); + if (newNodeMetadataOpt.isEmpty()) { + cf.complete(null); + return; + } + NodeMetadata newNodeMetadata = newNodeMetadataOpt.get(); + newNodeMetadata.setNodeEpoch(epoch); + this.nodeMetadata = newNodeMetadata; + AutomqRegisterNodeRequestData.TagCollection tagCollection = new AutomqRegisterNodeRequestData.TagCollection(); + newNodeMetadata.getTags().forEach((k, v) -> tagCollection.add(new AutomqRegisterNodeRequestData.Tag().setKey(k).setValue(v))); + AutomqRegisterNodeRequestData request = new AutomqRegisterNodeRequestData() + .setNodeId(nodeId) + .setNodeEpoch(epoch) + .setWalConfig(newNodeMetadata.getWalConfig()) + .setTags(tagCollection); + + WrapRequest req = new WrapRequest() { + @Override + public ApiKeys apiKey() { + return ApiKeys.AUTOMQ_REGISTER_NODE; + } + + @Override + public AbstractRequest.Builder toRequestBuilder() { + return new AutomqRegisterNodeRequest.Builder(request); + } + + @Override + public String toString() { + return request.toString(); + } + }; + + ControllerRequestSender.RequestTask task = new ControllerRequestSender.RequestTask<>(req, cf, + response -> { + AutomqRegisterNodeResponseData resp = response.data(); + Errors code = Errors.forCode(resp.errorCode()); + switch (code) { + case NONE: + return withSuccess(null); + case NODE_EPOCH_EXPIRED: + LOGGER.error("Node epoch expired: {}, code: {}", req, code); + throw code.exception(); + default: + LOGGER.error("Error while AUTOMQ_REGISTER_NODE: {}, code: {}, retry later", req, code); + return withRetry(); + } + }); + this.requestSender.send(task); + } catch (Throwable e) { + cf.completeExceptionally(e); + } + } + + @Override + public CompletableFuture getNodeMetadata() { + return lastCf.thenApply(nil -> this.nodeMetadata); + } + + private CompletableFuture> getNodeMetadata0() { + AutomqGetNodesRequestData request = new AutomqGetNodesRequestData().setNodeIds(List.of(nodeId)); + + WrapRequest req = new WrapRequest() { + @Override + public ApiKeys apiKey() { + return ApiKeys.AUTOMQ_GET_NODES; + } + + @Override + public AbstractRequest.Builder toRequestBuilder() { + return new AutomqGetNodesRequest.Builder(request); + } + + @Override + public String toString() { + return request.toString(); + } + }; + + CompletableFuture> future = new CompletableFuture<>(); + ControllerRequestSender.RequestTask> task = new ControllerRequestSender.RequestTask<>(req, future, + response -> { + AutomqGetNodesResponseData resp = response.data(); + Errors code = Errors.forCode(resp.errorCode()); + switch (code) { + case NONE: + return withSuccess(resp.nodes().stream().map(NodeManagerStub::from).filter(n -> n.getNodeId() == nodeId).findAny()); + default: + LOGGER.error("Error while AUTOMQ_GET_NODES: {}, code: {}, retry later", req, code); + return withRetry(); + } + }); + this.requestSender.send(task); + return future; + } + + static NodeMetadata from(AutomqGetNodesResponseData.NodeMetadata src) { + return new NodeMetadata( + src.nodeId(), + src.nodeEpoch(), + src.walConfig(), + src.tags().stream().collect(Collectors.toMap(AutomqGetNodesResponseData.Tag::key, AutomqGetNodesResponseData.Tag::value)) + ); + } +} diff --git a/core/src/main/scala/kafka/log/stream/s3/node/NoopNodeManager.java b/core/src/main/scala/kafka/log/stream/s3/node/NoopNodeManager.java new file mode 100644 index 0000000000..6e81842be0 --- /dev/null +++ b/core/src/main/scala/kafka/log/stream/s3/node/NoopNodeManager.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.stream.s3.node; + +import org.apache.kafka.controller.stream.NodeMetadata; + +import java.util.Collections; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +public class NoopNodeManager implements NodeManager { + private final int nodeId; + private final long nodeEpoch; + + public NoopNodeManager(int nodeId, long nodeEpoch) { + this.nodeId = nodeId; + this.nodeEpoch = nodeEpoch; + } + + @Override + public CompletableFuture update(Function> updater) { + return CompletableFuture.completedFuture(null); + } + + @Override + public CompletableFuture getNodeMetadata() { + return CompletableFuture.completedFuture(new NodeMetadata(nodeId, nodeEpoch, "", Collections.emptyMap())); + } +} diff --git a/core/src/main/scala/kafka/log/stream/s3/objects/ControllerObjectManager.java b/core/src/main/scala/kafka/log/stream/s3/objects/ControllerObjectManager.java index 8307d848c4..1fd8cdeea0 100644 --- a/core/src/main/scala/kafka/log/stream/s3/objects/ControllerObjectManager.java +++ b/core/src/main/scala/kafka/log/stream/s3/objects/ControllerObjectManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.objects; @@ -34,6 +42,8 @@ import com.automq.stream.s3.compact.CompactOperations; import com.automq.stream.s3.exceptions.AutoMQException; +import com.automq.stream.s3.exceptions.CompactedObjectsNotFoundException; +import com.automq.stream.s3.exceptions.ObjectNotCommittedException; import com.automq.stream.s3.metadata.S3ObjectMetadata; import com.automq.stream.s3.objects.CommitStreamSetObjectHook; import com.automq.stream.s3.objects.CommitStreamSetObjectRequest; @@ -180,7 +190,7 @@ public Builder toRequestBuilder() { throw Errors.forCode(resp.errorCode()).exception(); case OBJECT_NOT_EXIST: case COMPACTED_OBJECTS_NOT_FOUND: - throw code.exception(); + throw new CompactedObjectsNotFoundException(); default: LOGGER.error("Error while committing stream set object: {}, code: {}, retry later", request, code); return ResponseHandleResult.withRetry(); @@ -231,7 +241,9 @@ public Builder toRequestBuilder() { throw Errors.forCode(resp.errorCode()).exception(); case OBJECT_NOT_EXIST: case COMPACTED_OBJECTS_NOT_FOUND: - throw code.exception(); + throw new CompactedObjectsNotFoundException(); + case OBJECT_NOT_COMMITED: + throw new ObjectNotCommittedException(); case STREAM_NOT_EXIST: case STREAM_FENCED: LOGGER.warn("Stream fenced or not exist: {}, code: {}", request, Errors.forCode(resp.errorCode())); diff --git a/core/src/main/scala/kafka/log/stream/s3/objects/Convertor.java b/core/src/main/scala/kafka/log/stream/s3/objects/Convertor.java index f07850775b..7e81200222 100644 --- a/core/src/main/scala/kafka/log/stream/s3/objects/Convertor.java +++ b/core/src/main/scala/kafka/log/stream/s3/objects/Convertor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.objects; diff --git a/core/src/main/scala/kafka/log/stream/s3/streams/ControllerStreamManager.java b/core/src/main/scala/kafka/log/stream/s3/streams/ControllerStreamManager.java index 739c678d37..175693ba42 100644 --- a/core/src/main/scala/kafka/log/stream/s3/streams/ControllerStreamManager.java +++ b/core/src/main/scala/kafka/log/stream/s3/streams/ControllerStreamManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.streams; @@ -51,6 +59,7 @@ import com.automq.stream.s3.metadata.StreamState; import com.automq.stream.s3.streams.StreamCloseHook; import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.s3.streams.StreamMetadataListener; import com.automq.stream.utils.FutureUtil; import com.automq.stream.utils.LogContext; @@ -120,6 +129,7 @@ public String toString() { .map(m -> new StreamMetadata(m.streamId(), m.epoch(), m.startOffset(), m.endOffset(), StreamState.OPENED)) .collect(Collectors.toList())); case NODE_EPOCH_EXPIRED: + case NODE_FENCED: logger.error("Node epoch expired: {}, code: {}", req, code); throw code.exception(); default: @@ -136,6 +146,11 @@ public CompletableFuture> getStreams(List streamIds) return CompletableFuture.completedFuture(this.streamMetadataManager.getStreamMetadataList(streamIds)); } + @Override + public StreamMetadataListener.Handle addMetadataListener(long streamId, StreamMetadataListener listener) { + return streamMetadataManager.addMetadataListener(streamId, listener); + } + @Override public CompletableFuture createStream(Map tags) { CreateStreamRequest request = new CreateStreamRequest().setNodeId(nodeId); @@ -177,6 +192,7 @@ public Builder toRequestBuilder() { return ResponseHandleResult.withSuccess(resp.streamId()); case NODE_EPOCH_EXPIRED: case NODE_EPOCH_NOT_EXIST: + case NODE_FENCED: logger.error("Node epoch expired or not exist: {}, code: {}", req, Errors.forCode(resp.errorCode())); throw Errors.forCode(resp.errorCode()).exception(); default: @@ -238,6 +254,7 @@ public String toString() { new StreamMetadata(streamId, epoch, resp.startOffset(), resp.nextOffset(), StreamState.OPENED)); case NODE_EPOCH_EXPIRED: case NODE_EPOCH_NOT_EXIST: + case NODE_FENCED: logger.error("Node epoch expired or not exist, stream {}, epoch {}, code: {}", streamId, epoch, code); throw code.exception(); case STREAM_FENCED: @@ -250,6 +267,9 @@ public String toString() { case STREAM_NOT_CLOSED: logger.warn("open stream fail: {}, epoch {}, code: STREAM_NOT_CLOSED, retry later", streamId, epoch); return ResponseHandleResult.withRetry(); + case NODE_LOCKED: + logger.warn("[NODE_LOCKED]open stream fail: {}, epoch {}", streamId, epoch); + throw code.exception(); default: logger.error("Error while opening stream: {}, epoch {}, code: {}, retry later", streamId, epoch, code); return ResponseHandleResult.withRetry(); @@ -298,6 +318,7 @@ public Builder toRequestBuilder() { return ResponseHandleResult.withSuccess(null); case NODE_EPOCH_EXPIRED: case NODE_EPOCH_NOT_EXIST: + case NODE_FENCED: logger.error("Node epoch expired or not exist: {}, code: {}", request, Errors.forCode(resp.errorCode())); throw Errors.forCode(resp.errorCode()).exception(); case STREAM_NOT_EXIST: @@ -375,6 +396,7 @@ public Builder toRequestBuilder() { return ResponseHandleResult.withSuccess(null); case NODE_EPOCH_EXPIRED: case NODE_EPOCH_NOT_EXIST: + case NODE_FENCED: logger.error("Node epoch expired or not exist: {}, code: {}", request, Errors.forCode(resp.errorCode())); throw Errors.forCode(resp.errorCode()).exception(); case STREAM_NOT_EXIST: @@ -429,6 +451,7 @@ public Builder toRequestBuilder() { return ResponseHandleResult.withSuccess(null); case NODE_EPOCH_EXPIRED: case NODE_EPOCH_NOT_EXIST: + case NODE_FENCED: logger.error("Node epoch expired or not exist: {}, code: {}", request, Errors.forCode(resp.errorCode())); throw Errors.forCode(resp.errorCode()).exception(); case STREAM_NOT_EXIST: diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/ContextUtils.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/ContextUtils.java deleted file mode 100644 index edfa336050..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/ContextUtils.java +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry; - -import com.automq.stream.s3.context.AppendContext; -import com.automq.stream.s3.context.FetchContext; -import com.automq.stream.s3.trace.context.TraceContext; - -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; -import io.opentelemetry.sdk.OpenTelemetrySdk; - -public class ContextUtils { - public static FetchContext creaetFetchContext() { - return new FetchContext(createTraceContext()); - } - - public static AppendContext createAppendContext() { - return new AppendContext(createTraceContext()); - } - - public static TraceContext createTraceContext() { - OpenTelemetrySdk openTelemetrySdk = TelemetryManager.getOpenTelemetrySdk(); - boolean isTraceEnabled = openTelemetrySdk != null && TelemetryManager.isTraceEnable(); - Tracer tracer = null; - if (isTraceEnabled) { - tracer = openTelemetrySdk.getTracer(TelemetryConstants.TELEMETRY_SCOPE_NAME); - } - return new TraceContext(isTraceEnabled, tracer, Context.current()); - } - -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/MetricsConstants.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/MetricsConstants.java deleted file mode 100644 index 48739065f9..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/MetricsConstants.java +++ /dev/null @@ -1,21 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * Use of this software is governed by the Business Source License - * included in the file BSL.md - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry; - -public class MetricsConstants { - public static final String SERVICE_NAME = "service.name"; - public static final String SERVICE_INSTANCE = "service.instance.id"; - public static final String HOST_NAME = "host.name"; - public static final String INSTANCE = "instance"; - public static final String JOB = "job"; - public static final String NODE_TYPE = "node.type"; -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/TelemetryConstants.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/TelemetryConstants.java deleted file mode 100644 index 83875ac7c1..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/TelemetryConstants.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry; - -import io.opentelemetry.api.common.AttributeKey; - -public class TelemetryConstants { - // The maximum number of unique attribute combinations for a single metric - public static final int CARDINALITY_LIMIT = 20000; - public static final String COMMON_JMX_YAML_CONFIG_PATH = "/jmx/rules/common.yaml"; - public static final String BROKER_JMX_YAML_CONFIG_PATH = "/jmx/rules/broker.yaml"; - public static final String CONTROLLER_JMX_YAML_CONFIG_PATH = "/jmx/rules/controller.yaml"; - public static final String TELEMETRY_SCOPE_NAME = "automq_for_kafka"; - public static final String KAFKA_METRICS_PREFIX = "kafka_stream_"; - public static final String KAFKA_WAL_METRICS_PREFIX = "kafka_wal_"; - public static final AttributeKey STREAM_ID_NAME = AttributeKey.longKey("streamId"); - public static final AttributeKey START_OFFSET_NAME = AttributeKey.longKey("startOffset"); - public static final AttributeKey END_OFFSET_NAME = AttributeKey.longKey("endOffset"); - public static final AttributeKey MAX_BYTES_NAME = AttributeKey.longKey("maxBytes"); -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/TelemetryManager.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/TelemetryManager.java deleted file mode 100644 index bb0c68af30..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/TelemetryManager.java +++ /dev/null @@ -1,232 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry; - -import kafka.log.stream.s3.telemetry.exporter.MetricsExporter; -import kafka.log.stream.s3.telemetry.exporter.MetricsExporterURI; -import kafka.log.stream.s3.telemetry.otel.OTelHistogramReporter; -import kafka.server.KafkaConfig; - -import org.apache.kafka.server.ProcessRole; -import org.apache.kafka.server.metrics.KafkaYammerMetrics; -import org.apache.kafka.server.metrics.s3stream.S3StreamKafkaMetricsManager; - -import com.automq.stream.s3.metrics.MetricsConfig; -import com.automq.stream.s3.metrics.MetricsLevel; -import com.automq.stream.s3.metrics.S3StreamMetricsManager; -import com.automq.stream.s3.wal.metrics.ObjectWALMetricsManager; - -import org.apache.commons.lang3.StringUtils; -import org.apache.commons.lang3.tuple.Pair; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.slf4j.bridge.SLF4JBridgeHandler; - -import java.io.IOException; -import java.io.InputStream; -import java.net.InetAddress; -import java.util.ArrayList; -import java.util.List; -import java.util.Locale; - -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator; -import io.opentelemetry.api.common.Attributes; -import io.opentelemetry.api.common.AttributesBuilder; -import io.opentelemetry.api.metrics.Meter; -import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; -import io.opentelemetry.context.propagation.ContextPropagators; -import io.opentelemetry.context.propagation.TextMapPropagator; -import io.opentelemetry.instrumentation.jmx.engine.JmxMetricInsight; -import io.opentelemetry.instrumentation.jmx.engine.MetricConfiguration; -import io.opentelemetry.instrumentation.jmx.yaml.RuleParser; -import io.opentelemetry.instrumentation.runtimemetrics.java8.Cpu; -import io.opentelemetry.instrumentation.runtimemetrics.java8.GarbageCollector; -import io.opentelemetry.instrumentation.runtimemetrics.java8.MemoryPools; -import io.opentelemetry.instrumentation.runtimemetrics.java8.Threads; -import io.opentelemetry.sdk.OpenTelemetrySdk; -import io.opentelemetry.sdk.OpenTelemetrySdkBuilder; -import io.opentelemetry.sdk.metrics.SdkMeterProvider; -import io.opentelemetry.sdk.metrics.SdkMeterProviderBuilder; -import io.opentelemetry.sdk.metrics.export.MetricReader; -import io.opentelemetry.sdk.metrics.internal.SdkMeterProviderUtil; -import io.opentelemetry.sdk.resources.Resource; -import scala.collection.immutable.Set; - -public class TelemetryManager { - private static final Logger LOGGER = LoggerFactory.getLogger(TelemetryManager.class); - private final KafkaConfig kafkaConfig; - private final String clusterId; - protected final List metricReaderList; - private final List autoCloseableList; - private final OTelHistogramReporter oTelHistogramReporter; - private JmxMetricInsight jmxMetricInsight; - private OpenTelemetrySdk openTelemetrySdk; - - public TelemetryManager(KafkaConfig kafkaConfig, String clusterId) { - this.kafkaConfig = kafkaConfig; - this.clusterId = clusterId; - this.metricReaderList = new ArrayList<>(); - this.autoCloseableList = new ArrayList<>(); - this.oTelHistogramReporter = new OTelHistogramReporter(KafkaYammerMetrics.defaultRegistry()); - // redirect JUL from OpenTelemetry SDK to SLF4J - SLF4JBridgeHandler.removeHandlersForRootLogger(); - SLF4JBridgeHandler.install(); - } - - private String getHostName() { - try { - return InetAddress.getLocalHost().getHostName(); - } catch (Exception e) { - LOGGER.error("Failed to get host name", e); - return "unknown"; - } - } - - public void init() { - OpenTelemetrySdkBuilder openTelemetrySdkBuilder = OpenTelemetrySdk.builder(); - openTelemetrySdkBuilder.setMeterProvider(buildMeterProvider(kafkaConfig)); - openTelemetrySdk = openTelemetrySdkBuilder - .setPropagators(ContextPropagators.create(TextMapPropagator.composite( - W3CTraceContextPropagator.getInstance(), W3CBaggagePropagator.getInstance()))) - .build(); - - addJmxMetrics(openTelemetrySdk); - addJvmMetrics(openTelemetrySdk); - - // initialize S3Stream metrics - Meter meter = openTelemetrySdk.getMeter(TelemetryConstants.TELEMETRY_SCOPE_NAME); - initializeMetricsManager(meter); - } - - protected SdkMeterProvider buildMeterProvider(KafkaConfig kafkaConfig) { - AttributesBuilder baseAttributesBuilder = Attributes.builder() - .put(MetricsConstants.SERVICE_NAME, clusterId) - .put(MetricsConstants.SERVICE_INSTANCE, String.valueOf(kafkaConfig.nodeId())) - .put(MetricsConstants.HOST_NAME, getHostName()) - .put(MetricsConstants.JOB, clusterId) // for Prometheus HTTP server compatibility - .put(MetricsConstants.INSTANCE, String.valueOf(kafkaConfig.nodeId())); // for Aliyun Prometheus compatibility - List> extraAttributes = kafkaConfig.automq().baseLabels(); - if (extraAttributes != null) { - for (Pair pair : extraAttributes) { - baseAttributesBuilder.put(pair.getKey(), pair.getValue()); - } - } - - Resource resource = Resource.empty().toBuilder() - .putAll(baseAttributesBuilder.build()) - .build(); - SdkMeterProviderBuilder sdkMeterProviderBuilder = SdkMeterProvider.builder().setResource(resource); - MetricsExporterURI metricsExporterURI = buildMetricsExporterURI(clusterId, kafkaConfig); - if (metricsExporterURI != null) { - for (MetricsExporter metricsExporter : metricsExporterURI.metricsExporters()) { - MetricReader metricReader = metricsExporter.asMetricReader(); - metricReaderList.add(metricReader); - SdkMeterProviderUtil.registerMetricReaderWithCardinalitySelector(sdkMeterProviderBuilder, metricReader, - instrumentType -> TelemetryConstants.CARDINALITY_LIMIT); - } - } - return sdkMeterProviderBuilder.build(); - } - - protected MetricsExporterURI buildMetricsExporterURI(String clusterId, KafkaConfig kafkaConfig) { - return MetricsExporterURI.parse(clusterId, kafkaConfig); - } - - protected void initializeMetricsManager(Meter meter) { - S3StreamMetricsManager.configure(new MetricsConfig(metricsLevel(), Attributes.empty(), kafkaConfig.s3ExporterReportIntervalMs())); - S3StreamMetricsManager.initMetrics(meter, TelemetryConstants.KAFKA_METRICS_PREFIX); - - S3StreamKafkaMetricsManager.configure(new MetricsConfig(metricsLevel(), Attributes.empty(), kafkaConfig.s3ExporterReportIntervalMs())); - S3StreamKafkaMetricsManager.initMetrics(meter, TelemetryConstants.KAFKA_METRICS_PREFIX); - - // kraft controller may not have s3WALPath config. - ObjectWALMetricsManager.initMetrics(meter, TelemetryConstants.KAFKA_WAL_METRICS_PREFIX); - this.oTelHistogramReporter.start(meter); - } - - private void addJmxMetrics(OpenTelemetry ot) { - jmxMetricInsight = JmxMetricInsight.createService(ot, kafkaConfig.s3ExporterReportIntervalMs()); - MetricConfiguration conf = new MetricConfiguration(); - - Set roles = kafkaConfig.processRoles(); - buildMetricConfiguration(conf, TelemetryConstants.COMMON_JMX_YAML_CONFIG_PATH); - if (roles.contains(ProcessRole.BrokerRole)) { - buildMetricConfiguration(conf, TelemetryConstants.BROKER_JMX_YAML_CONFIG_PATH); - } - if (roles.contains(ProcessRole.ControllerRole)) { - buildMetricConfiguration(conf, TelemetryConstants.CONTROLLER_JMX_YAML_CONFIG_PATH); - } - jmxMetricInsight.start(conf); - } - - private void buildMetricConfiguration(MetricConfiguration conf, String path) { - try (InputStream ins = this.getClass().getResourceAsStream(path)) { - RuleParser parser = RuleParser.get(); - parser.addMetricDefsTo(conf, ins, path); - } catch (Exception e) { - LOGGER.error("Failed to parse JMX config file: {}", path, e); - } - } - - private void addJvmMetrics(OpenTelemetry openTelemetry) { - // JVM metrics - autoCloseableList.addAll(MemoryPools.registerObservers(openTelemetry)); - autoCloseableList.addAll(Cpu.registerObservers(openTelemetry)); - autoCloseableList.addAll(GarbageCollector.registerObservers(openTelemetry)); - autoCloseableList.addAll(Threads.registerObservers(openTelemetry)); - } - - protected MetricsLevel metricsLevel() { - String levelStr = kafkaConfig.s3MetricsLevel(); - if (StringUtils.isBlank(levelStr)) { - return MetricsLevel.INFO; - } - try { - String up = levelStr.toUpperCase(Locale.ENGLISH); - return MetricsLevel.valueOf(up); - } catch (Exception e) { - LOGGER.error("illegal metrics level: {}", levelStr); - return MetricsLevel.INFO; - } - } - - public void shutdown() { - autoCloseableList.forEach(autoCloseable -> { - try { - autoCloseable.close(); - } catch (Exception e) { - LOGGER.error("Failed to close auto closeable", e); - } - }); - metricReaderList.forEach(metricReader -> { - metricReader.forceFlush(); - try { - metricReader.close(); - } catch (IOException e) { - LOGGER.error("Failed to close metric reader", e); - } - }); - if (openTelemetrySdk != null) { - openTelemetrySdk.close(); - } - } - - // Deprecated methods, leave for compatibility - public static boolean isTraceEnable() { - return false; - } - - public static OpenTelemetrySdk getOpenTelemetrySdk() { - return null; - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/ExporterConstants.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/ExporterConstants.java deleted file mode 100644 index bb9d3f9503..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/ExporterConstants.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -public class ExporterConstants { - public static final String OTLP_TYPE = "otlp"; - public static final String PROMETHEUS_TYPE = "prometheus"; - public static final String OPS_TYPE = "ops"; - public static final String URI_DELIMITER = "://?"; - public static final String ENDPOINT = "endpoint"; - public static final String PROTOCOL = "protocol"; - public static final String COMPRESSION = "compression"; - public static final String HOST = "host"; - public static final String PORT = "port"; - public static final String COMPRESSION_GZIP = "gzip"; - public static final String COMPRESSION_NONE = "none"; - public static final String OTLP_GRPC_PROTOCOL = "grpc"; - public static final String OTLP_HTTP_PROTOCOL = "http"; - public static final String DEFAULT_PROM_HOST = "localhost"; - public static final int DEFAULT_PROM_PORT = 9090; - public static final int DEFAULT_EXPORTER_TIMEOUT_MS = 30000; -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporter.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporter.java deleted file mode 100644 index 3f1a55c4f5..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporter.java +++ /dev/null @@ -1,18 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -import io.opentelemetry.sdk.metrics.export.MetricReader; - -public interface MetricsExporter { - MetricReader asMetricReader(); -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporterType.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporterType.java deleted file mode 100644 index 23861777f9..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporterType.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -public enum MetricsExporterType { - OTLP("otlp"), - PROMETHEUS("prometheus"), - OPS("ops"); - - private final String type; - - MetricsExporterType(String type) { - this.type = type; - } - - public String getType() { - return type; - } - - public static MetricsExporterType fromString(String type) { - for (MetricsExporterType exporterType : MetricsExporterType.values()) { - if (exporterType.getType().equalsIgnoreCase(type)) { - return exporterType; - } - } - throw new IllegalArgumentException("Invalid metrics exporter type: " + type); - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporterURI.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporterURI.java deleted file mode 100644 index 558d34b522..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/MetricsExporterURI.java +++ /dev/null @@ -1,116 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -import kafka.server.KafkaConfig; - -import org.apache.kafka.common.utils.Utils; - -import com.automq.stream.s3.operator.BucketURI; -import com.automq.stream.utils.URIUtils; - -import org.apache.commons.lang3.tuple.Pair; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.net.URI; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.Map; - -import software.amazon.awssdk.annotations.NotNull; - -public class MetricsExporterURI { - private static final Logger LOGGER = LoggerFactory.getLogger(MetricsExporterURI.class); - private final List metricsExporters; - - public MetricsExporterURI(List metricsExporters) { - this.metricsExporters = metricsExporters == null ? new ArrayList<>() : metricsExporters; - } - - public static MetricsExporter parseExporter(String clusterId, KafkaConfig kafkaConfig, String uriStr) { - try { - URI uri = new URI(uriStr); - String type = uri.getScheme(); - if (Utils.isBlank(type)) { - LOGGER.error("Invalid metrics exporter URI: {}, exporter type is missing", uriStr); - return null; - } - Map> queries = URIUtils.splitQuery(uri); - return parseExporter(clusterId, kafkaConfig, type, queries); - } catch (Exception e) { - LOGGER.warn("Parse metrics exporter URI {} failed", uriStr, e); - return null; - } - } - - public static MetricsExporter parseExporter(String clusterId, KafkaConfig kafkaConfig, String type, Map> queries) { - MetricsExporterType exporterType = MetricsExporterType.fromString(type); - switch (exporterType) { - case OTLP: - return buildOTLPExporter(kafkaConfig.s3ExporterReportIntervalMs(), queries); - case PROMETHEUS: - return buildPrometheusExporter(queries, kafkaConfig.automq().baseLabels()); - case OPS: - return buildOpsExporter(clusterId, kafkaConfig.nodeId(), kafkaConfig.s3ExporterReportIntervalMs(), - kafkaConfig.automq().opsBuckets(), kafkaConfig.automq().baseLabels()); - default: - return null; - } - } - - public static @NotNull MetricsExporterURI parse(String clusterId, KafkaConfig kafkaConfig) { - String uriStr = kafkaConfig.automq().metricsExporterURI(); - if (Utils.isBlank(uriStr)) { - return new MetricsExporterURI(Collections.emptyList()); - } - String[] exporterUri = uriStr.split(","); - if (exporterUri.length == 0) { - return new MetricsExporterURI(Collections.emptyList()); - } - List exporters = new ArrayList<>(); - for (String uri : exporterUri) { - if (Utils.isBlank(uri)) { - continue; - } - MetricsExporter exporter = parseExporter(clusterId, kafkaConfig, uri); - if (exporter != null) { - exporters.add(exporter); - } - } - return new MetricsExporterURI(exporters); - } - - public static MetricsExporter buildOTLPExporter(int intervalMs, Map> queries) { - String endpoint = URIUtils.getString(queries, ExporterConstants.ENDPOINT, ""); - String protocol = URIUtils.getString(queries, ExporterConstants.PROTOCOL, OTLPProtocol.GRPC.getProtocol()); - String compression = URIUtils.getString(queries, ExporterConstants.COMPRESSION, OTLPCompressionType.NONE.getType()); - return new OTLPMetricsExporter(intervalMs, endpoint, protocol, compression); - } - - public static MetricsExporter buildPrometheusExporter(Map> queries, List> baseLabels) { - String host = URIUtils.getString(queries, ExporterConstants.HOST, ExporterConstants.DEFAULT_PROM_HOST); - int port = Integer.parseInt(URIUtils.getString(queries, ExporterConstants.PORT, String.valueOf(ExporterConstants.DEFAULT_PROM_PORT))); - return new PrometheusMetricsExporter(host, port, baseLabels); - } - - public static MetricsExporter buildOpsExporter(String clusterId, int nodeId, int intervalMs, List opsBuckets, - List> baseLabels) { - return new OpsMetricsExporter(clusterId, nodeId, intervalMs, opsBuckets, baseLabels); - } - - public List metricsExporters() { - return metricsExporters; - } - -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPCompressionType.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPCompressionType.java deleted file mode 100644 index 481da55b59..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPCompressionType.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -public enum OTLPCompressionType { - GZIP("gzip"), - NONE("none"); - - private final String type; - - OTLPCompressionType(String type) { - this.type = type; - } - - public String getType() { - return type; - } - - public static OTLPCompressionType fromString(String type) { - for (OTLPCompressionType compressionType : OTLPCompressionType.values()) { - if (compressionType.getType().equalsIgnoreCase(type)) { - return compressionType; - } - } - throw new IllegalArgumentException("Invalid OTLP compression type: " + type); - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPProtocol.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPProtocol.java deleted file mode 100644 index 305f9b417a..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OTLPProtocol.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -public enum OTLPProtocol { - GRPC("grpc"), - HTTP("http"); - - private final String protocol; - - OTLPProtocol(String protocol) { - this.protocol = protocol; - } - - public String getProtocol() { - return protocol; - } - - public static OTLPProtocol fromString(String protocol) { - for (OTLPProtocol otlpProtocol : OTLPProtocol.values()) { - if (otlpProtocol.getProtocol().equalsIgnoreCase(protocol)) { - return otlpProtocol; - } - } - throw new IllegalArgumentException("Invalid OTLP protocol: " + protocol); - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OpsMetricsExporter.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OpsMetricsExporter.java deleted file mode 100644 index 9bd9f94e34..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/OpsMetricsExporter.java +++ /dev/null @@ -1,106 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -import kafka.server.KafkaRaftServer; - -import com.automq.shell.AutoMQApplication; -import com.automq.shell.metrics.S3MetricsConfig; -import com.automq.shell.metrics.S3MetricsExporter; -import com.automq.stream.s3.operator.BucketURI; -import com.automq.stream.s3.operator.ObjectStorage; -import com.automq.stream.s3.operator.ObjectStorageFactory; - -import org.apache.commons.lang3.tuple.Pair; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.time.Duration; -import java.util.List; - -import io.opentelemetry.sdk.metrics.export.MetricReader; -import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader; - -public class OpsMetricsExporter implements MetricsExporter { - private static final Logger LOGGER = LoggerFactory.getLogger(OpsMetricsExporter.class); - private final String clusterId; - private final int nodeId; - private final int intervalMs; - private final List opsBuckets; - private final List> baseLabels; - - public OpsMetricsExporter(String clusterId, int nodeId, int intervalMs, List opsBuckets, List> baseLabels) { - if (opsBuckets == null || opsBuckets.isEmpty()) { - throw new IllegalArgumentException("At least one bucket URI must be provided for ops metrics exporter"); - } - this.clusterId = clusterId; - this.nodeId = nodeId; - this.intervalMs = intervalMs; - this.opsBuckets = opsBuckets; - this.baseLabels = baseLabels; - LOGGER.info("OpsMetricsExporter initialized with clusterId: {}, nodeId: {}, intervalMs: {}, opsBuckets: {}", - clusterId, nodeId, intervalMs, opsBuckets); - } - - public String clusterId() { - return clusterId; - } - - public int nodeId() { - return nodeId; - } - - public int intervalMs() { - return intervalMs; - } - - public List opsBuckets() { - return opsBuckets; - } - - @Override - public MetricReader asMetricReader() { - BucketURI bucket = opsBuckets.get(0); - ObjectStorage objectStorage = ObjectStorageFactory.instance().builder(bucket).threadPrefix("ops-metric").build(); - S3MetricsConfig metricsConfig = new S3MetricsConfig() { - @Override - public String clusterId() { - return clusterId; - } - - @Override - public boolean isActiveController() { - KafkaRaftServer raftServer = AutoMQApplication.getBean(KafkaRaftServer.class); - return raftServer != null && raftServer.controller().exists(controller -> controller.controller() != null - && controller.controller().isActive()); - } - - @Override - public int nodeId() { - return nodeId; - } - - @Override - public ObjectStorage objectStorage() { - return objectStorage; - } - - @Override - public List> baseLabels() { - return baseLabels; - } - }; - S3MetricsExporter s3MetricsExporter = new S3MetricsExporter(metricsConfig); - s3MetricsExporter.start(); - return PeriodicMetricReader.builder(s3MetricsExporter).setInterval(Duration.ofMillis(intervalMs)).build(); - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/PrometheusMetricsExporter.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/PrometheusMetricsExporter.java deleted file mode 100644 index 0c5bd2d8e2..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/exporter/PrometheusMetricsExporter.java +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -import kafka.log.stream.s3.telemetry.MetricsConstants; - -import org.apache.kafka.common.utils.Utils; - -import org.apache.commons.lang3.tuple.Pair; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.util.List; -import java.util.Set; -import java.util.stream.Collectors; - -import io.opentelemetry.exporter.prometheus.PrometheusHttpServer; -import io.opentelemetry.sdk.metrics.export.MetricReader; - -public class PrometheusMetricsExporter implements MetricsExporter { - private static final Logger LOGGER = LoggerFactory.getLogger(PrometheusMetricsExporter.class); - private final String host; - private final int port; - private final Set baseLabelKeys; - - public PrometheusMetricsExporter(String host, int port, List> baseLabels) { - if (Utils.isBlank(host)) { - throw new IllegalArgumentException("Illegal Prometheus host"); - } - if (port <= 0) { - throw new IllegalArgumentException("Illegal Prometheus port"); - } - this.host = host; - this.port = port; - this.baseLabelKeys = baseLabels.stream().map(Pair::getKey).collect(Collectors.toSet()); - LOGGER.info("PrometheusMetricsExporter initialized with host: {}, port: {}", host, port); - } - - public String host() { - return host; - } - - public int port() { - return port; - } - - @Override - public MetricReader asMetricReader() { - return PrometheusHttpServer.builder() - .setHost(host) - .setPort(port) - .setAllowedResourceAttributesFilter(resourceAttributes -> - MetricsConstants.JOB.equals(resourceAttributes) - || MetricsConstants.INSTANCE.equals(resourceAttributes) - || MetricsConstants.HOST_NAME.equals(resourceAttributes) - || baseLabelKeys.contains(resourceAttributes)) - .build(); - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelHistogramReporter.java b/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelHistogramReporter.java deleted file mode 100644 index b2c97170e2..0000000000 --- a/core/src/main/scala/kafka/log/stream/s3/telemetry/otel/OTelHistogramReporter.java +++ /dev/null @@ -1,66 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.otel; - -import com.yammer.metrics.core.Metric; -import com.yammer.metrics.core.MetricName; -import com.yammer.metrics.core.MetricsRegistry; -import com.yammer.metrics.core.MetricsRegistryListener; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import io.opentelemetry.api.metrics.Meter; - -// This class is responsible for transforming yammer histogram metrics (mean, max) into OTel metrics -public class OTelHistogramReporter implements MetricsRegistryListener { - private static final Logger LOGGER = LoggerFactory.getLogger(OTelHistogramReporter.class); - private final MetricsRegistry metricsRegistry; - private final OTelMetricsProcessor metricsProcessor; - private volatile Meter meter; - - public OTelHistogramReporter(MetricsRegistry metricsRegistry) { - this.metricsRegistry = metricsRegistry; - this.metricsProcessor = new OTelMetricsProcessor(); - } - - public void start(Meter meter) { - this.meter = meter; - this.metricsProcessor.init(meter); - metricsRegistry.addListener(this); - LOGGER.info("OTelHistogramReporter started"); - } - - @Override - public void onMetricAdded(MetricName name, Metric metric) { - if (OTelMetricUtils.isInterestedMetric(name)) { - if (this.meter == null) { - LOGGER.info("Not initialized yet, skipping metric: {}", name); - return; - } - try { - metric.processWith(this.metricsProcessor, name, null); - } catch (Throwable t) { - LOGGER.error("Failed to process metric: {}", name, t); - } - } - } - - @Override - public void onMetricRemoved(MetricName name) { - try { - this.metricsProcessor.remove(name); - } catch (Throwable ignored) { - - } - } -} diff --git a/core/src/main/scala/kafka/log/stream/s3/wal/ConfirmWal.java b/core/src/main/scala/kafka/log/stream/s3/wal/ConfirmWal.java new file mode 100644 index 0000000000..c5e7f61b78 --- /dev/null +++ b/core/src/main/scala/kafka/log/stream/s3/wal/ConfirmWal.java @@ -0,0 +1,210 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.stream.s3.wal; + +import kafka.log.stream.s3.node.NodeManager; + +import org.apache.kafka.common.utils.ThreadUtils; +import org.apache.kafka.controller.stream.NodeMetadata; + +import com.automq.stream.s3.exceptions.AutoMQException; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.trace.context.TraceContext; +import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.RecoverResult; +import com.automq.stream.s3.wal.WalFactory; +import com.automq.stream.s3.wal.WalFactory.BuildOptions; +import com.automq.stream.s3.wal.WalHandle; +import com.automq.stream.s3.wal.WriteAheadLog; +import com.automq.stream.s3.wal.common.WALMetadata; +import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.utils.IdURI; +import com.automq.stream.utils.LogContext; +import com.automq.stream.utils.Threads; + +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; + +import java.io.IOException; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; + +public class ConfirmWal implements WriteAheadLog { + private final Logger logger; + private final int nodeId; + private final long nodeEpoch; + private final String walConfigs; + private final boolean failoverMode; + + private final WalFactory factory; + private final NodeManager nodeManager; + private final WalHandle walHandle; + private final ExecutorService executor; + private volatile WriteAheadLog wal; + private String currentWalConfigs; + + public ConfirmWal(int nodeId, long nodeEpoch, String walConfigs, boolean failoverMode, + WalFactory factory, NodeManager nodeManager, WalHandle walHandle) { + String name = String.format("CONFIRM_WAL-%s-%s%s", nodeId, nodeEpoch, failoverMode ? "-F" : ""); + this.logger = new LogContext("[" + name + "] ").logger(ConfirmWal.class); + this.nodeId = nodeId; + this.nodeEpoch = nodeEpoch; + this.walConfigs = walConfigs; + this.failoverMode = failoverMode; + this.factory = factory; + this.nodeManager = nodeManager; + this.walHandle = walHandle; + this.executor = Threads.newFixedThreadPoolWithMonitor(2, name, true, logger); + + try { + // Init register node config if the node is the first time to start. + NodeMetadata oldNodeMetadata = this.nodeManager.getNodeMetadata().get(); + if (StringUtils.isBlank(oldNodeMetadata.getWalConfig())) { + // https://github.com/AutoMQ/automq/releases/tag/1.5.0 + // https://github.com/AutoMQ/automq/pull/2517 + // AutoMQ supports registering wal config to kraft after version 1.5.0. + // So we need to recover the data even if the old wal config is empty. + currentWalConfigs = walConfigs; + logger.info("The WAL is the 'first time' to start."); + this.wal = buildRecoverWal(currentWalConfigs, nodeEpoch - 1).get(); + } else { + if (nodeEpoch < oldNodeMetadata.getNodeEpoch()) { + throw new AutoMQException("The node epoch is less than the current node epoch: " + nodeEpoch + " < " + oldNodeMetadata.getNodeEpoch()); + } + logger.info("Using the old config {} for recovering", oldNodeMetadata); + currentWalConfigs = oldNodeMetadata.getWalConfig(); + // We should use the new nodeEpoch here. + // Consider the case: the node re-bootstraps between new wal start and register to NodeManager in reset. + // The wal epoch is already set to the new nodeEpoch. + // So we need to use the new nodeEpoch to recover the data. + this.wal = buildRecoverWal(oldNodeMetadata.getWalConfig(), nodeEpoch).get(); + } + + + } catch (Throwable e) { + throw new AutoMQException(e); + } + } + + @Override + public WriteAheadLog start() throws IOException { + return wal.start(); + } + + @Override + public void shutdownGracefully() { + wal.shutdownGracefully(); + ThreadUtils.shutdownExecutorServiceQuietly(executor, 10, TimeUnit.SECONDS); + } + + @Override + public WALMetadata metadata() { + return wal.metadata(); + } + + @Override + public String uri() { + return walConfigs; + } + + @Override + public CompletableFuture append(TraceContext context, StreamRecordBatch streamRecordBatch) throws OverCapacityException { + return wal.append(context, streamRecordBatch); + } + + @Override + public CompletableFuture get(RecordOffset recordOffset) { + return wal.get(recordOffset); + } + + @Override + public CompletableFuture> get(RecordOffset startOffset, RecordOffset endOffset) { + return wal.get(startOffset, endOffset); + } + + @Override + public RecordOffset confirmOffset() { + return wal.confirmOffset(); + } + + @Override + public Iterator recover() { + return wal.recover(); + } + + @Override + public CompletableFuture reset() { + return CompletableFuture.runAsync(() -> { + try { + wal.reset().get(); + wal.shutdownGracefully(); + if (failoverMode) { + releasePermission(currentWalConfigs).get(); + return; + } + this.currentWalConfigs = walConfigs; + this.wal = buildWal(currentWalConfigs).get(); + wal.start(); + Iterator it = wal.recover(); + logger.info("Register new WAL configs: {}", currentWalConfigs); + nodeManager.updateWal(currentWalConfigs).join(); + if (it.hasNext()) { + throw new AutoMQException("[WARN] The WAL 'should be' empty, try reboot to recover"); + } + wal.reset().get(); + } catch (Throwable e) { + logger.error("Reset WAL failed:", e); + throw new AutoMQException(e); + } + }, executor); + } + + @Override + public CompletableFuture trim(RecordOffset offset) { + return wal.trim(offset); + } + + private CompletableFuture buildRecoverWal(String kraftWalConfigs, long nodeEpoch) { + IdURI uri = IdURI.parse(kraftWalConfigs); + CompletableFuture cf = walHandle + .acquirePermission(nodeId, nodeEpoch, uri, new WalHandle.AcquirePermissionOptions().failoverMode(failoverMode)); + return cf.thenApplyAsync(nil -> factory.build(uri, BuildOptions.builder().nodeEpoch(nodeEpoch).openMode(failoverMode ? OpenMode.FAILOVER : OpenMode.READ_WRITE).build()), executor); + } + + private CompletableFuture buildWal(String kraftWalConfigs) { + IdURI uri = IdURI.parse(kraftWalConfigs); + WalHandle.AcquirePermissionOptions options = new WalHandle.AcquirePermissionOptions() + .timeoutMs(Long.MAX_VALUE) + .failoverMode(false); + CompletableFuture cf = walHandle + .acquirePermission(nodeId, nodeEpoch, uri, options); + return cf.thenApplyAsync(nil -> factory.build(uri, BuildOptions.builder().nodeEpoch(nodeEpoch).openMode(OpenMode.READ_WRITE).build()), executor); + } + + private CompletableFuture releasePermission(String kraftWalConfigs) { + IdURI uri = IdURI.parse(kraftWalConfigs); + return walHandle.releasePermission(uri, new WalHandle.ReleasePermissionOptions()); + } +} diff --git a/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java b/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java new file mode 100644 index 0000000000..b0e1d626c8 --- /dev/null +++ b/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java @@ -0,0 +1,82 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.stream.s3.wal; + +import com.automq.shell.AutoMQApplication; +import com.automq.stream.s3.network.NetworkBandwidthLimiter; +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; +import com.automq.stream.s3.wal.ReservationService; +import com.automq.stream.s3.wal.WalFactory; +import com.automq.stream.s3.wal.WriteAheadLog; +import com.automq.stream.s3.wal.impl.object.ObjectReservationService; +import com.automq.stream.s3.wal.impl.object.ObjectWALConfig; +import com.automq.stream.s3.wal.impl.object.ObjectWALService; +import com.automq.stream.utils.IdURI; +import com.automq.stream.utils.Time; + +import java.util.Locale; +import java.util.Map; + +public class DefaultWalFactory implements WalFactory { + private final int nodeId; + private final Map objectTagging; + private final NetworkBandwidthLimiter networkInboundLimiter; + private final NetworkBandwidthLimiter networkOutboundLimiter; + + public DefaultWalFactory(int nodeId, Map objectTagging, + NetworkBandwidthLimiter networkInboundLimiter, NetworkBandwidthLimiter networkOutboundLimiter) { + this.nodeId = nodeId; + this.objectTagging = objectTagging; + this.networkInboundLimiter = networkInboundLimiter; + this.networkOutboundLimiter = networkOutboundLimiter; + } + + @Override + public WriteAheadLog build(IdURI uri, BuildOptions options) { + //noinspection SwitchStatementWithTooFewBranches + switch (uri.protocol().toUpperCase(Locale.ENGLISH)) { + case "S3": + BucketURI bucketURI = to(uri); + ObjectStorage walObjectStorage = ObjectStorageFactory.instance() + .builder(bucketURI) + .tagging(objectTagging) + .inboundLimiter(networkInboundLimiter) + .outboundLimiter(networkOutboundLimiter) + .build(); + + ObjectWALConfig.Builder configBuilder = ObjectWALConfig.builder().withURI(uri) + .withClusterId(AutoMQApplication.getClusterId()) + .withNodeId(nodeId) + .withEpoch(options.nodeEpoch()) + .withOpenMode(options.openMode()); + ReservationService reservationService = new ObjectReservationService(AutoMQApplication.getClusterId(), walObjectStorage, walObjectStorage.bucketId()); + configBuilder.withReservationService(reservationService); + return new ObjectWALService(Time.SYSTEM, walObjectStorage, configBuilder.build()); + default: + throw new IllegalArgumentException("Unsupported WAL protocol: " + uri.protocol()); + } + } + + static BucketURI to(IdURI uri) { + return BucketURI.parse(uri.encode()); + } +} diff --git a/core/src/main/scala/kafka/log/streamaspect/CachedLogSegments.java b/core/src/main/scala/kafka/log/streamaspect/CachedLogSegments.java index 0c1bcac581..12e081da85 100644 --- a/core/src/main/scala/kafka/log/streamaspect/CachedLogSegments.java +++ b/core/src/main/scala/kafka/log/streamaspect/CachedLogSegments.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -41,7 +49,7 @@ public LogSegment add(LogSegment segment) { public void remove(long offset) { synchronized (this) { if (null != activeSegment && offset == activeSegment.baseOffset()) { - activeSegment = null; + activeSegment = super.lastSegment().orElse(null); } super.remove(offset); } diff --git a/core/src/main/scala/kafka/log/streamaspect/AlwaysSuccessClient.java b/core/src/main/scala/kafka/log/streamaspect/ClientWrapper.java similarity index 66% rename from core/src/main/scala/kafka/log/streamaspect/AlwaysSuccessClient.java rename to core/src/main/scala/kafka/log/streamaspect/ClientWrapper.java index 4ebbfabaae..dc63e654c3 100644 --- a/core/src/main/scala/kafka/log/streamaspect/AlwaysSuccessClient.java +++ b/core/src/main/scala/kafka/log/streamaspect/ClientWrapper.java @@ -1,17 +1,28 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.errors.s3.NodeEpochExpiredException; +import org.apache.kafka.common.errors.s3.NodeEpochNotExistException; +import org.apache.kafka.common.errors.s3.NodeFencedException; import org.apache.kafka.common.utils.ThreadUtils; import com.automq.stream.api.AppendResult; @@ -49,29 +60,29 @@ import static com.automq.stream.utils.FutureUtil.cause; -public class AlwaysSuccessClient implements Client { +public class ClientWrapper implements Client { - private static final Logger LOGGER = LoggerFactory.getLogger(AlwaysSuccessClient.class); + private static final Logger LOGGER = LoggerFactory.getLogger(ClientWrapper.class); public static final Set HALT_ERROR_CODES = Set.of( - ErrorCode.EXPIRED_STREAM_EPOCH, - ErrorCode.STREAM_ALREADY_CLOSED, - ErrorCode.OFFSET_OUT_OF_RANGE_BOUNDS + ErrorCode.EXPIRED_STREAM_EPOCH, + ErrorCode.STREAM_ALREADY_CLOSED, + ErrorCode.OFFSET_OUT_OF_RANGE_BOUNDS ); private final ScheduledExecutorService streamManagerRetryScheduler = Executors.newScheduledThreadPool(1, - ThreadUtils.createThreadFactory("stream-manager-retry-%d", true)); + ThreadUtils.createThreadFactory("stream-manager-retry-scheduler-%d", true)); private final ExecutorService streamManagerCallbackExecutors = Executors.newFixedThreadPool(1, - ThreadUtils.createThreadFactory("stream-manager-callback-executor-%d", true)); + ThreadUtils.createThreadFactory("stream-manager-callback-executor-%d", true)); private final ScheduledExecutorService generalRetryScheduler = Executors.newScheduledThreadPool(1, - ThreadUtils.createThreadFactory("general-retry-scheduler-%d", true)); + ThreadUtils.createThreadFactory("general-retry-scheduler-%d", true)); private final ExecutorService generalCallbackExecutors = Executors.newFixedThreadPool(4, - ThreadUtils.createThreadFactory("general-callback-scheduler-%d", true)); + ThreadUtils.createThreadFactory("general-callback-executor-%d", true)); private final Client innerClient; private volatile StreamClient streamClient; private final HashedWheelTimer fetchTimeout = new HashedWheelTimer( - ThreadUtils.createThreadFactory("fetch-timeout-%d", true), - 1, TimeUnit.SECONDS, 512); + ThreadUtils.createThreadFactory("fetch-timeout-%d", true), + 1, TimeUnit.SECONDS, 512); - public AlwaysSuccessClient(Client client) { + public ClientWrapper(Client client) { this.innerClient = client; } @@ -111,6 +122,7 @@ public void shutdown() { streamManagerCallbackExecutors.shutdownNow(); generalRetryScheduler.shutdownNow(); generalCallbackExecutors.shutdownNow(); + fetchTimeout.stop(); } /** @@ -148,6 +160,22 @@ private static boolean maybeHaltAndCompleteWaitingFuture(Throwable t, Completabl return true; } + private static CompletableFuture failureHandle(CompletableFuture cf) { + return cf.whenComplete((rst, ex) -> { + if (ex != null) { + ex = FutureUtil.cause(ex); + if (ex instanceof NodeEpochExpiredException + || ex instanceof NodeEpochNotExistException + || ex instanceof NodeFencedException) { + LOGGER.error("The node is fenced, force shutdown the node", ex); + //noinspection CallToPrintStackTrace + ex.printStackTrace(); + Runtime.getRuntime().halt(1); + } + } + }); + } + private class StreamClientImpl implements StreamClient { private final StreamClient streamClient; @@ -158,29 +186,12 @@ public StreamClientImpl(StreamClient streamClient) { @Override public CompletableFuture createAndOpenStream(CreateStreamOptions options) { - CompletableFuture cf = new CompletableFuture<>(); - createAndOpenStream0(options, cf); - return cf; - } - - private void createAndOpenStream0(CreateStreamOptions options, CompletableFuture cf) { - streamClient.createAndOpenStream(options).whenCompleteAsync((stream, ex) -> { - FutureUtil.suppress(() -> { - if (ex != null) { - LOGGER.error("Create and open stream fail, retry later", ex); - streamManagerRetryScheduler.schedule(() -> createAndOpenStream0(options, cf), 3, TimeUnit.SECONDS); - } else { - cf.complete(new StreamImpl(stream)); - } - }, LOGGER); - }, streamManagerCallbackExecutors); + return failureHandle(streamClient.createAndOpenStream(options).thenApplyAsync(rst -> rst, streamManagerCallbackExecutors)); } @Override public CompletableFuture openStream(long streamId, OpenStreamOptions options) { - CompletableFuture cf = new CompletableFuture<>(); - openStream0(streamId, options, cf); - return cf; + return failureHandle(streamClient.openStream(streamId, options).thenApplyAsync(rst -> rst, streamManagerCallbackExecutors)); } @Override @@ -191,21 +202,6 @@ public Optional getStream(long streamId) { public void shutdown() { streamClient.shutdown(); } - - private void openStream0(long streamId, OpenStreamOptions options, CompletableFuture cf) { - streamClient.openStream(streamId, options).whenCompleteAsync((stream, ex) -> { - FutureUtil.suppress(() -> { - if (ex != null) { - if (!maybeHaltAndCompleteWaitingFuture(ex, cf)) { - LOGGER.error("Open stream[{}]({}) fail, retry later", streamId, options.epoch(), ex); - streamManagerRetryScheduler.schedule(() -> openStream0(streamId, options, cf), 3, TimeUnit.SECONDS); - } - } else { - cf.complete(new StreamImpl(stream)); - } - }, LOGGER); - }, generalCallbackExecutors); - } } private class StreamImpl implements Stream { @@ -235,6 +231,11 @@ public long confirmOffset() { return stream.confirmOffset(); } + @Override + public void confirmOffset(long offset) { + stream.confirmOffset(offset); + } + @Override public long nextOffset() { return stream.nextOffset(); @@ -250,7 +251,8 @@ public CompletableFuture append(AppendContext context, RecordBatch } @Override - public CompletableFuture fetch(FetchContext context, long startOffset, long endOffset, int maxBytesHint) { + public CompletableFuture fetch(FetchContext context, long startOffset, long endOffset, + int maxBytesHint) { CompletableFuture cf = new CompletableFuture<>(); Timeout timeout = fetchTimeout.newTimeout(t -> LOGGER.warn("fetch timeout, stream[{}] [{}, {})", streamId(), startOffset, endOffset), 1, TimeUnit.MINUTES); stream.fetch(context, startOffset, endOffset, maxBytesHint).whenComplete((rst, e) -> FutureUtil.suppress(() -> { @@ -292,42 +294,17 @@ private void trim0(long newStartOffset, CompletableFuture cf) { @Override public CompletableFuture close() { - CompletableFuture cf = new CompletableFuture<>(); - close0(cf); - return cf; - } - - private void close0(CompletableFuture cf) { - stream.close().whenCompleteAsync((rst, ex) -> FutureUtil.suppress(() -> { - if (ex != null) { - if (!maybeHaltAndCompleteWaitingFuture(ex, cf)) { - LOGGER.error("Close stream[{}] failed, retry later", streamId(), ex); - generalRetryScheduler.schedule(() -> close0(cf), 3, TimeUnit.SECONDS); - } - } else { - cf.complete(rst); - } - }, LOGGER), generalCallbackExecutors); + return failureHandle(stream.close().thenApplyAsync(nil -> nil, streamManagerCallbackExecutors)); } @Override public CompletableFuture destroy() { - CompletableFuture cf = new CompletableFuture<>(); - destroy0(cf); - return cf; + return failureHandle(stream.destroy().thenApplyAsync(nil -> nil, streamManagerCallbackExecutors)); } - private void destroy0(CompletableFuture cf) { - stream.destroy().whenCompleteAsync((rst, ex) -> FutureUtil.suppress(() -> { - if (ex != null) { - if (!maybeHaltAndCompleteWaitingFuture(ex, cf)) { - LOGGER.error("Destroy stream[{}] failed, retry later", streamId(), ex); - generalRetryScheduler.schedule(() -> destroy0(cf), 3, TimeUnit.SECONDS); - } - } else { - cf.complete(rst); - } - }, LOGGER), generalCallbackExecutors); + @Override + public CompletableFuture lastAppendFuture() { + return stream.lastAppendFuture(); } } } diff --git a/core/src/main/scala/kafka/log/streamaspect/DefaultElasticStreamSlice.java b/core/src/main/scala/kafka/log/streamaspect/DefaultElasticStreamSlice.java index 9d2fbfc3b8..8d70b2b75f 100644 --- a/core/src/main/scala/kafka/log/streamaspect/DefaultElasticStreamSlice.java +++ b/core/src/main/scala/kafka/log/streamaspect/DefaultElasticStreamSlice.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -34,10 +42,8 @@ public class DefaultElasticStreamSlice implements ElasticStreamSlice { */ private final long startOffsetInStream; private final Stream stream; - /** - * next relative offset to be appended to this segment. - */ - private long nextOffset; + // The relative endOffset of sealed stream slice + private long endOffset = Offsets.NOOP_OFFSET; private boolean sealed = false; public DefaultElasticStreamSlice(Stream stream, SliceRange sliceRange) { @@ -47,15 +53,13 @@ public DefaultElasticStreamSlice(Stream stream, SliceRange sliceRange) { // new stream slice this.startOffsetInStream = streamNextOffset; sliceRange.start(startOffsetInStream); - this.nextOffset = 0L; } else if (sliceRange.end() == Offsets.NOOP_OFFSET) { // unsealed stream slice this.startOffsetInStream = sliceRange.start(); - this.nextOffset = streamNextOffset - startOffsetInStream; } else { // sealed stream slice this.startOffsetInStream = sliceRange.start(); - this.nextOffset = sliceRange.end() - startOffsetInStream; + this.endOffset = sliceRange.end() - startOffsetInStream; this.sealed = true; } } @@ -65,7 +69,6 @@ public CompletableFuture append(AppendContext context, RecordBatch if (sealed) { return FutureUtil.failedFuture(new IllegalStateException("stream segment " + this + " is sealed")); } - nextOffset += recordBatch.count(); return stream.append(context, recordBatch).thenApply(AppendResultWrapper::new); } @@ -78,23 +81,18 @@ public CompletableFuture fetch(FetchContext context, long startOffs @Override public long nextOffset() { - return nextOffset; + return endOffset != Offsets.NOOP_OFFSET ? endOffset : (stream.nextOffset() - startOffsetInStream); } @Override public long confirmOffset() { - return stream.confirmOffset() - startOffsetInStream; - } - - @Override - public long startOffsetInStream() { - return startOffsetInStream; + return endOffset != Offsets.NOOP_OFFSET ? endOffset : (stream.confirmOffset() - startOffsetInStream); } @Override public SliceRange sliceRange() { if (sealed) { - return SliceRange.of(startOffsetInStream, startOffsetInStream + nextOffset); + return SliceRange.of(startOffsetInStream, startOffsetInStream + endOffset); } else { return SliceRange.of(startOffsetInStream, Offsets.NOOP_OFFSET); } @@ -102,7 +100,10 @@ public SliceRange sliceRange() { @Override public void seal() { - this.sealed = true; + if (!sealed) { + sealed = true; + endOffset = stream.nextOffset() - startOffsetInStream; + } } @Override @@ -113,9 +114,9 @@ public Stream stream() { @Override public String toString() { return "DefaultElasticStreamSlice{" + - "startOffsetInStream=" + startOffsetInStream + - ", stream=[id=" + stream.streamId() + ", startOffset=" + stream.startOffset() + ", nextOffset=" + stream.nextOffset() + "]" + - ", nextOffset=" + nextOffset + + ", streamId=" + stream.streamId() + + ", slice=" + sliceRange() + + ", nextOffset=" + nextOffset() + ", sealed=" + sealed + '}'; } diff --git a/core/src/main/scala/kafka/log/streamaspect/DefaultOpenStreamChecker.java b/core/src/main/scala/kafka/log/streamaspect/DefaultOpenStreamChecker.java index d4043d4020..49a59c5547 100644 --- a/core/src/main/scala/kafka/log/streamaspect/DefaultOpenStreamChecker.java +++ b/core/src/main/scala/kafka/log/streamaspect/DefaultOpenStreamChecker.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -22,19 +30,23 @@ import com.automq.stream.s3.metadata.StreamState; +import java.util.Arrays; + public class DefaultOpenStreamChecker implements OpenStreamChecker { + private final int nodeId; private final KRaftMetadataCache metadataCache; - public DefaultOpenStreamChecker(KRaftMetadataCache metadataCache) { + public DefaultOpenStreamChecker(int nodeId, KRaftMetadataCache metadataCache) { + this.nodeId = nodeId; this.metadataCache = metadataCache; } @Override public boolean check(Uuid topicId, int partition, long streamId, long epoch) throws StreamFencedException { - return metadataCache.safeRun(image -> DefaultOpenStreamChecker.check(image, topicId, partition, streamId, epoch)); + return metadataCache.safeRun(image -> DefaultOpenStreamChecker.check(image, topicId, partition, streamId, epoch, nodeId)); } - public static boolean check(MetadataImage image, Uuid topicId, int partition, long streamId, long epoch) throws StreamFencedException { + public static boolean check(MetadataImage image, Uuid topicId, int partition, long streamId, long epoch, int currentNodeId) throws StreamFencedException { // When ABA reassign happens: // 1. Assign P0 to broker0 with epoch=0, broker0 opens the partition // 2. Assign P0 to broker1 with epoch=1, broker1 waits for the partition to be closed @@ -52,6 +64,10 @@ public static boolean check(MetadataImage image, Uuid topicId, int partition, lo if (currentEpoch > epoch) { throw new StreamFencedException(String.format("partition=%s-%d with epoch=%d is fenced by new leader epoch=%d", topicId, partition, epoch, currentEpoch)); } + if (!contains(partitionImage.isr, currentNodeId)) { + throw new StreamFencedException(String.format("partition=%s-%d with epoch=%d move to other nodes %s", topicId, partition, epoch, Arrays.toString(partitionImage.isr))); + } + S3StreamMetadataImage stream = image.streamsMetadata().getStreamMetadata(streamId); if (stream == null) { throw new StreamFencedException(String.format("streamId=%d cannot be found, it may be deleted or not created yet", streamId)); @@ -60,4 +76,16 @@ public static boolean check(MetadataImage image, Uuid topicId, int partition, lo throw new StreamFencedException(String.format("streamId=%d with epoch=%d is fenced by new leader epoch=%d", streamId, epoch, stream.getEpoch())); return StreamState.CLOSED.equals(stream.state()); } + + private static boolean contains(int[] isr, int nodeId) { + if (isr == null) { + return false; + } + for (int replica : isr) { + if (replica == nodeId) { + return true; + } + } + return false; + } } diff --git a/core/src/main/scala/kafka/log/streamaspect/DefaultStreamSliceSupplier.java b/core/src/main/scala/kafka/log/streamaspect/DefaultStreamSliceSupplier.java index 7c8a22bace..8ad60dbd3e 100644 --- a/core/src/main/scala/kafka/log/streamaspect/DefaultStreamSliceSupplier.java +++ b/core/src/main/scala/kafka/log/streamaspect/DefaultStreamSliceSupplier.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticCheckoutPointFileWithHandler.scala b/core/src/main/scala/kafka/log/streamaspect/ElasticCheckoutPointFileWithHandler.scala index 072ec2d869..4b761294a9 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticCheckoutPointFileWithHandler.scala +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticCheckoutPointFileWithHandler.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpoint.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpoint.java index 82b94b4e0b..3813ea2318 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpoint.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpoint.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ @@ -46,4 +54,4 @@ public synchronized void writeIfDirExists(Collection epochs) { public synchronized List read() { return meta.entries(); } -} \ No newline at end of file +} diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpointMeta.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpointMeta.java index 81432576d7..b1a3fc8cc0 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpointMeta.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLeaderEpochCheckpointMeta.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLog.scala b/core/src/main/scala/kafka/log/streamaspect/ElasticLog.scala index 1934a2b788..a03d4419d9 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLog.scala +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLog.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect @@ -14,6 +22,7 @@ package kafka.log.streamaspect import com.automq.stream.api.{Client, CreateStreamOptions, KeyValue, OpenStreamOptions} import com.automq.stream.utils.{FutureUtil, Systems} import io.netty.buffer.Unpooled +import kafka.cluster.PartitionSnapshot import kafka.log.LocalLog.CleanedFileSuffix import kafka.log._ import kafka.log.streamaspect.ElasticLogFileRecords.{BatchIteratorRecordsAdaptor, PooledMemoryRecords} @@ -36,7 +45,7 @@ import java.io.{File, IOException} import java.nio.ByteBuffer import java.nio.file.Files import java.util -import java.util.Optional +import java.util.{Collections, Optional} import java.util.concurrent._ import java.util.concurrent.atomic.{AtomicLong, AtomicReference} import scala.collection.mutable.ListBuffer @@ -87,7 +96,8 @@ class ElasticLog(val metaStream: MetaStream, topicPartition: TopicPartition, logDirFailureChannel: LogDirFailureChannel, val _initStartOffset: Long = 0, - leaderEpoch: Long + leaderEpoch: Long, + snapshotRead: Boolean = false ) extends LocalLog(__dir, _config, segments, partitionMeta.getRecoverOffset, _nextOffsetMetadata, scheduler, time, topicPartition, logDirFailureChannel) { import ElasticLog._ @@ -100,7 +110,7 @@ class ElasticLog(val metaStream: MetaStream, var confirmOffsetChangeListener: Option[() => Unit] = None private val appendAckQueue = new LinkedBlockingQueue[Long]() - private val appendAckThread = APPEND_CALLBACK_EXECUTOR(math.abs(logIdent.hashCode % APPEND_CALLBACK_EXECUTOR.length)) + val appendAckThread = APPEND_CALLBACK_EXECUTOR(math.abs(logIdent.hashCode % APPEND_CALLBACK_EXECUTOR.length)) @volatile private[log] var lastAppendAckFuture: Future[?] = CompletableFuture.completedFuture(null) private val readAsyncThread = READ_ASYNC_EXECUTOR(math.abs(logIdent.hashCode % READ_ASYNC_EXECUTOR.length)) @@ -110,6 +120,7 @@ class ElasticLog(val metaStream: MetaStream, streamManager.setListener((_, event) => { if (event == ElasticStreamMetaEvent.STREAM_DO_CREATE) { logSegmentManager.asyncPersistLogMeta() + logSegmentManager.notifySegmentUpdate(); } }) @@ -168,10 +179,16 @@ class ElasticLog(val metaStream: MetaStream, } private def persistLogMeta(): Unit = { + if (snapshotRead) { + return + } logSegmentManager.persistLogMeta() } private def persistPartitionMeta(): Unit = { + if (snapshotRead) { + return + } persistMeta(metaStream, MetaKeyValue.of(MetaStream.PARTITION_META_KEY, ElasticPartitionMeta.encode(partitionMeta))) if (isDebugEnabled) { debug(s"${logIdent}save partition meta $partitionMeta") @@ -263,13 +280,7 @@ class ElasticLog(val metaStream: MetaStream, } private[log] def confirmOffset: LogOffsetMetadata = { - val confirmOffset = _confirmOffset.get() - val offsetUpperBound = logSegmentManager.offsetUpperBound.get() - if (offsetUpperBound != null && offsetUpperBound.messageOffset < confirmOffset.messageOffset) { - offsetUpperBound - } else { - confirmOffset - } + _confirmOffset.get() } override private[log] def flush(offset: Long): Unit = { @@ -296,9 +307,11 @@ class ElasticLog(val metaStream: MetaStream, maxOffsetMetadata: LogOffsetMetadata, includeAbortedTxns: Boolean): CompletableFuture[FetchDataInfo] = { maybeHandleIOExceptionAsync(s"Exception while reading from $topicPartition in dir ${dir.getParent}") { - trace(s"Reading maximum $maxLength bytes at offset $startOffset from log with " + - s"total length ${segments.sizeInBytes} bytes") - // get LEO from super class + if (isTraceEnabled) { + trace(s"Reading maximum $maxLength bytes at offset $startOffset from log with " + + s"total length ${segments.sizeInBytes} bytes") + } + // get LEO from super class val endOffsetMetadata = nextOffsetMetadata val endOffset = endOffsetMetadata.messageOffset val segmentOpt = segments.lastSegment @@ -478,7 +491,11 @@ class ElasticLog(val metaStream: MetaStream, * Directly close all streams of the log. */ def closeStreams(): CompletableFuture[Void] = { - CompletableFuture.allOf(streamManager.close(), metaStream.close()) + if (snapshotRead) { + CompletableFuture.allOf(streamManager.close()) + } else { + CompletableFuture.allOf(streamManager.close(), metaStream.close()) + } } def updateLogStartOffset(offset: Long): Unit = { @@ -579,10 +596,62 @@ class ElasticLog(val metaStream: MetaStream, newSegment } } + + override private[log] def truncateFullyAndStartAt(newOffset: Long): Iterable[LogSegment] = { + val rst = super.truncateFullyAndStartAt(newOffset) + _confirmOffset.set(logEndOffsetMetadata) + rst + } + + def snapshot(snapshot: PartitionSnapshot.Builder): Unit = { + snapshot.logMeta(logSegmentManager.logMeta()) + snapshot.logEndOffset(logEndOffsetMetadata) + logSegmentManager.streams().forEach(stream => { + snapshot.streamEndOffset(stream.streamId(), stream.nextOffset()) + snapshot.addStreamLastAppendFuture(stream.lastAppendFuture()); + }) + val lastSegmentOpt = segments.lastSegment() + if (lastSegmentOpt.isPresent) { + snapshot.lastTimestampOffset(lastSegmentOpt.get().asInstanceOf[ElasticLogSegment].timeIndex().lastEntry()) + } + } + + def snapshot(snapshot: PartitionSnapshot): Unit = { + val logMeta = snapshot.logMeta() + if (logMeta != null && !logMeta.getSegmentMetas.isEmpty) { + logMeta.getStreamMap.forEach((name, streamId) => { + streamManager.createIfNotExist(name, streamId) + }) + segments.clear() + logMeta.getSegmentMetas.forEach(segMeta => { + val segment = new ElasticLogSegment(dir, segMeta, streamSliceManager, config, time, (_, _) => {}, logIdent) + segments.add(segment) + }) + } + var logEndOffset = snapshot.logEndOffset() + val segmentBaseOffset = segments.floorSegment(logEndOffset.messageOffset).get().baseOffset() + logEndOffset = new LogOffsetMetadata(logEndOffset.messageOffset, segmentBaseOffset, logEndOffset.relativePositionInSegment); + + streamManager.streams().forEach((_, stream) => { + val endOffset = snapshot.streamEndOffsets().get(stream.streamId()) + if (endOffset != null) { + stream.confirmOffset(endOffset) + } + }) + val lastSegment = segments.lastSegment() + if (lastSegment.isPresent) { + lastSegment.get().asInstanceOf[ElasticLogSegment].snapshot(snapshot) + } + nextOffsetMetadata = logEndOffset + _confirmOffset.set(logEndOffset) + } } object ElasticLog extends Logging { - private val APPEND_PERMIT = 100 * 1024 * 1024 + private val APPEND_PERMIT = Systems.getEnvInt("AUTOMQ_APPEND_PERMIT_SIZE", + // autoscale the append permit size based on heap size, min 100MiB, max 1GiB, every 6GB heap add 100MiB permit + Math.min(1024, 100 * Math.max(1, (Systems.HEAP_MEMORY_SIZE / (1024 * 1024 * 1024) / 6)).asInstanceOf[Int]) * 1024 * 1024 + ) private val APPEND_PERMIT_SEMAPHORE = new Semaphore(APPEND_PERMIT) S3StreamKafkaMetricsManager.setLogAppendPermitNumSupplier(() => APPEND_PERMIT_SEMAPHORE.availablePermits()) @@ -623,7 +692,8 @@ object ElasticLog extends Logging { producerStateManagerConfig: ProducerStateManagerConfig, topicId: Option[Uuid], leaderEpoch: Long, - openStreamChecker: OpenStreamChecker + openStreamChecker: OpenStreamChecker, + snapshotRead: Boolean = false ): ElasticLog = { // TODO: better error mark for elastic log logDirFailureChannel.clearOfflineLogDirRecord(dir.getPath) @@ -647,6 +717,18 @@ object ElasticLog extends Logging { streamTags.put(StreamTags.Partition.KEY, StreamTags.Partition.encode(topicPartition.partition())) try { + if (snapshotRead) { + val logStreamManager = new ElasticLogStreamManager(Collections.emptyMap(), client.streamClient(), replicationFactor, leaderEpoch, streamTags, true) + val streamSliceManager = new ElasticStreamSliceManager(logStreamManager) + val segments = new CachedLogSegments(topicPartition) + partitionMeta = new ElasticPartitionMeta() + val leaderEpochCheckpointMeta = new ElasticLeaderEpochCheckpointMeta(LeaderEpochCheckpointFile.CURRENT_VERSION, new util.ArrayList[EpochEntry]()) + val producerStateManager = new ElasticProducerStateManager(topicPartition, dir, + maxTransactionTimeoutMs, producerStateManagerConfig, time, new util.TreeMap[java.lang.Long, ByteBuffer](), _ => CompletableFuture.completedFuture(null)) + return new ElasticLog(null, logStreamManager, streamSliceManager, producerStateManager, null, partitionMeta, leaderEpochCheckpointMeta, dir, config, + segments, new LogOffsetMetadata(0), scheduler, time, topicPartition, logDirFailureChannel, 0, leaderEpoch, true) + } + metaStream = if (metaNotExists) { val stream = createMetaStream(client, key, replicationFactor, leaderEpoch, streamTags, logIdent = logIdent) info(s"${logIdent}created a new meta stream: streamId=${stream.streamId()}") @@ -687,7 +769,7 @@ object ElasticLog extends Logging { maxTransactionTimeoutMs, producerStateManagerConfig, time, snapshotsMap, kv => metaStream.append(kv).thenApply(_ => null)) val logMeta: ElasticLogMeta = metaMap.get(MetaStream.LOG_META_KEY).map(m => m.asInstanceOf[ElasticLogMeta]).getOrElse(new ElasticLogMeta()) - logStreamManager = new ElasticLogStreamManager(logMeta.getStreamMap, client.streamClient(), replicationFactor, leaderEpoch, streamTags) + logStreamManager = new ElasticLogStreamManager(logMeta.getStreamMap, client.streamClient(), replicationFactor, leaderEpoch, streamTags, false) val streamSliceManager = new ElasticStreamSliceManager(logStreamManager) val logSegmentManager = new ElasticLogSegmentManager(metaStream, logStreamManager, logIdent) diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java index 4c94af6096..8d9d9affaa 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogFileRecords.java @@ -1,18 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; -import kafka.log.stream.s3.telemetry.ContextUtils; -import kafka.log.stream.s3.telemetry.TelemetryConstants; +import kafka.automq.zerozone.LinkRecord; +import kafka.automq.zerozone.ZeroZoneThreadLocalContext; import org.apache.kafka.common.network.TransferableChannel; import org.apache.kafka.common.record.AbstractRecords; @@ -36,8 +44,8 @@ import com.automq.stream.s3.ByteBufAlloc; import com.automq.stream.s3.context.AppendContext; import com.automq.stream.s3.context.FetchContext; -import com.automq.stream.s3.trace.TraceUtils; import com.automq.stream.utils.FutureUtil; +import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -52,11 +60,9 @@ import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; -import io.opentelemetry.api.common.Attributes; import static com.automq.stream.s3.ByteBufAlloc.POOLED_MEMORY_RECORDS; import static com.automq.stream.utils.FutureUtil.suppress; @@ -70,8 +76,6 @@ public class ElasticLogFileRecords implements AutoCloseable { private final ElasticStreamSlice streamSlice; // This is The base offset of the corresponding segment. private final long baseOffset; - private final AtomicLong nextOffset; - private final AtomicLong committedOffset; // Inflight append result. private volatile CompletableFuture lastAppend; private volatile ElasticResourceStatus status; @@ -80,16 +84,13 @@ public class ElasticLogFileRecords implements AutoCloseable { public ElasticLogFileRecords(ElasticStreamSlice streamSlice, long baseOffset, int size) { this.baseOffset = baseOffset; this.streamSlice = streamSlice; - long nextOffset = streamSlice.nextOffset(); // Note that size is generally used to // 1) show the physical size of a segment. In these cases, size is referred to decide whether to roll a new // segment, or calculate the cleaned size in a cleaning task, etc. If size is not correctly recorded for any // reason, the worst thing will be just a bigger segment than configured. // 2) show whether this segment is empty, i.e., size == 0. // Therefore, it is fine to use the nextOffset as a backoff value. - this.size = new AtomicInteger(size == 0 ? (int) nextOffset : size); - this.nextOffset = new AtomicLong(baseOffset + nextOffset); - this.committedOffset = new AtomicLong(baseOffset + nextOffset); + this.size = new AtomicInteger(size == 0 ? (int) Math.min(streamSlice.nextOffset(), Integer.MAX_VALUE / 2) : size); this.lastAppend = CompletableFuture.completedFuture(null); batches = batchesFrom(baseOffset); @@ -100,12 +101,16 @@ public int sizeInBytes() { return size.get(); } + public void size(int size) { + this.size.set(size); + } + public long nextOffset() { - return nextOffset.get(); + return baseOffset + streamSlice.nextOffset(); } public long appendedOffset() { - return nextOffset.get() - baseOffset; + return nextOffset() - baseOffset; } public CompletableFuture read(long startOffset, long maxOffset, int maxSize) { @@ -114,21 +119,15 @@ public CompletableFuture read(long startOffset, long maxOffset, int max } if (ReadHint.isReadAll()) { ReadOptions readOptions = ReadOptions.builder().fastRead(ReadHint.isFastRead()).pooledBuf(true).build(); - FetchContext fetchContext = ContextUtils.creaetFetchContext(); + FetchContext fetchContext = new FetchContext(); fetchContext.setReadOptions(readOptions); - Attributes attributes = Attributes.builder() - .put(TelemetryConstants.START_OFFSET_NAME, startOffset) - .put(TelemetryConstants.END_OFFSET_NAME, maxOffset) - .put(TelemetryConstants.MAX_BYTES_NAME, maxSize) - .build(); try { - return TraceUtils.runWithSpanAsync(fetchContext, attributes, "ElasticLogFileRecords::read", - () -> readAll0(fetchContext, startOffset, maxOffset, maxSize)); + return readAll0(fetchContext, startOffset, maxOffset, maxSize); } catch (Throwable ex) { return CompletableFuture.failedFuture(ex); } } else { - long endOffset = Utils.min(this.committedOffset.get(), maxOffset); + long endOffset = Utils.min(confirmOffset(), maxOffset); return CompletableFuture.completedFuture(new BatchIteratorRecordsAdaptor(this, startOffset, endOffset, maxSize)); } } @@ -136,21 +135,39 @@ public CompletableFuture read(long startOffset, long maxOffset, int max private CompletableFuture readAll0(FetchContext context, long startOffset, long maxOffset, int maxSize) { // calculate the relative offset in the segment, which may start from 0. long nextFetchOffset = startOffset - baseOffset; - long endOffset = Utils.min(this.committedOffset.get(), maxOffset) - baseOffset; + long endOffset = Utils.min(confirmOffset(), maxOffset) - baseOffset; if (nextFetchOffset >= endOffset) { return CompletableFuture.completedFuture(MemoryRecords.EMPTY); } - return fetch0(context, nextFetchOffset, endOffset, maxSize) - .thenApply(rst -> PooledMemoryRecords.of(baseOffset, rst, context.readOptions().pooledBuf())); + List results = new LinkedList<>(); + return fetch0(context, nextFetchOffset, endOffset, maxSize, results) + .whenComplete((nil, e) -> { + if (e != null) { + results.forEach(FetchResult::free); + results.clear(); + } + }) + .thenApply(nil -> PooledMemoryRecords.of(baseOffset, results, context.readOptions().pooledBuf())); } - private CompletableFuture> fetch0(FetchContext context, long startOffset, long endOffset, int maxSize) { + /** + * Fetch records from the {@link ElasticStreamSlice} + * + * @param context fetch context + * @param startOffset start offset + * @param endOffset end offset + * @param maxSize max size of the fetched records + * @param results result list to be filled + * @return a future that completes when reaching the end offset or the max size + */ + private CompletableFuture fetch0(FetchContext context, long startOffset, long endOffset, int maxSize, List results) { if (startOffset >= endOffset || maxSize <= 0) { - return CompletableFuture.completedFuture(new LinkedList<>()); + return CompletableFuture.completedFuture(null); } int adjustedMaxSize = Math.min(maxSize, 1024 * 1024); return streamSlice.fetch(context, startOffset, endOffset, adjustedMaxSize) .thenCompose(rst -> { + results.add(rst); long nextFetchOffset = startOffset; int readSize = 0; for (RecordBatchWithContext recordBatchWithContext : rst.recordBatchList()) { @@ -163,12 +180,10 @@ private CompletableFuture> fetch0(FetchContext context, } readSize += recordBatchWithContext.rawPayload().remaining(); } - return fetch0(context, nextFetchOffset, endOffset, maxSize - readSize) - .thenApply(rstList -> { - // add to first since we need to reverse the order. - rstList.addFirst(rst); - return rstList; - }); + if (readSize == 0) { + return CompletableFuture.completedFuture(null); + } + return fetch0(context, nextFetchOffset, endOffset, maxSize - readSize, results); }); } @@ -191,24 +206,26 @@ public int append(MemoryRecords records, long lastOffset) throws IOException { " bytes is too large for segment with current file position at " + size.get()); int appendSize = records.sizeInBytes(); // Note that the calculation of count requires strong consistency between nextOffset and the baseOffset of records. - int count = (int) (lastOffset - nextOffset.get()); + int count = (int) (lastOffset - nextOffset()); com.automq.stream.DefaultRecordBatch batch = new com.automq.stream.DefaultRecordBatch(count, 0, Collections.emptyMap(), records.buffer()); - - AppendContext context = ContextUtils.createAppendContext(); + AppendContext context = new AppendContext(); + ZeroZoneThreadLocalContext.WriteContext writeContext = ZeroZoneThreadLocalContext.writeContext(); + ByteBuf linkRecord = LinkRecord.encode(writeContext.channelOffset(), records); + if (linkRecord != null) { + context.linkRecord(linkRecord); + } CompletableFuture cf; try { - cf = TraceUtils.runWithSpanAsync(context, Attributes.empty(), "ElasticLogFileRecords::append", - () -> streamSlice.append(context, batch)); + cf = streamSlice.append(context, batch); } catch (Throwable ex) { throw new IOException("Failed to append to stream " + streamSlice.stream().streamId(), ex); + } finally { + writeContext.reset(); } - nextOffset.set(lastOffset); size.getAndAdd(appendSize); cf.whenComplete((rst, e) -> { - if (e == null) { - updateCommittedOffset(lastOffset); - } else if (e instanceof IOException) { + if (e instanceof IOException) { status = ElasticResourceStatus.FENCED; LOGGER.error("ElasticLogFileRecords[stream={}, baseOffset={}] fencing with ex: {}", streamSlice.stream().streamId(), baseOffset, e.getMessage()); } @@ -217,15 +234,8 @@ public int append(MemoryRecords records, long lastOffset) throws IOException { return appendSize; } - private void updateCommittedOffset(long newCommittedOffset) { - while (true) { - long oldCommittedOffset = this.committedOffset.get(); - if (oldCommittedOffset >= newCommittedOffset) { - break; - } else if (this.committedOffset.compareAndSet(oldCommittedOffset, newCommittedOffset)) { - break; - } - } + private long confirmOffset() { + return baseOffset + streamSlice.confirmOffset(); } public void flush() throws IOException { @@ -398,7 +408,8 @@ public long lastOffset() { } static class StreamSegmentInputStream implements LogInputStream { - private static final int FETCH_BATCH_SIZE = 64 * 1024; + @VisibleForTesting + protected static final int FETCH_BATCH_SIZE = 512 * 1024; private final ElasticLogFileRecords elasticLogFileRecords; private final Queue remaining = new LinkedList<>(); private final int maxSize; @@ -412,7 +423,7 @@ public StreamSegmentInputStream(FetchContext fetchContext, ElasticLogFileRecords this.elasticLogFileRecords = elasticLogFileRecords; this.maxSize = maxSize; this.nextFetchOffset = startOffset - elasticLogFileRecords.baseOffset; - this.endOffset = Utils.min(elasticLogFileRecords.committedOffset.get(), maxOffset) - elasticLogFileRecords.baseOffset; + this.endOffset = Utils.min(elasticLogFileRecords.confirmOffset(), maxOffset) - elasticLogFileRecords.baseOffset; } @@ -438,15 +449,15 @@ public RecordBatch nextBatch() throws IOException { buf = heapBuf; } readSize += buf.remaining(); + nextFetchOffset = Math.max(streamRecord.lastOffset(), nextFetchOffset); for (RecordBatch r : MemoryRecords.readableRecords(buf).batches()) { remaining.offer(r); - nextFetchOffset = r.lastOffset() - elasticLogFileRecords.baseOffset + 1; } } catch (Throwable e) { ElasticStreamSlice slice = elasticLogFileRecords.streamSlice; byte[] bytes = new byte[streamRecord.rawPayload().remaining()]; streamRecord.rawPayload().get(bytes); - LOGGER.error("next batch parse error, stream={} baseOffset={} payload={}", slice.stream().streamId(), slice.sliceRange().start() + streamRecord.baseOffset(), bytes); + LOGGER.error("next batch parse error, stream={} baseOffset={} payload={}", slice.stream().streamId(), slice.sliceRange().start() + streamRecord.baseOffset(), bytes, e); throw new RuntimeException(e); } }); diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogLoader.scala b/core/src/main/scala/kafka/log/streamaspect/ElasticLogLoader.scala index bc1929483a..15e6a4498c 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogLoader.scala +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogLoader.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogManager.scala b/core/src/main/scala/kafka/log/streamaspect/ElasticLogManager.scala index 2e2205acf5..12e67f9793 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogManager.scala +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogManager.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect @@ -35,7 +43,7 @@ class ElasticLogManager(val client: Client, val openStreamChecker: OpenStreamChe this.logIdent = s"[ElasticLogManager] " private val elasticLogs = new ConcurrentHashMap[TopicPartition, ElasticUnifiedLog]() - def getOrCreateLog(dir: File, + def createLog(dir: File, config: LogConfig, scheduler: Scheduler, time: Time, @@ -47,10 +55,10 @@ class ElasticLogManager(val client: Client, val openStreamChecker: OpenStreamChe topicId: Option[Uuid], leaderEpoch: Long = 0, ): ElasticUnifiedLog = { + val snapshotRead = OpenHint.isSnapshotRead val topicPartition = UnifiedLog.parseTopicPartitionName(dir) - val log = elasticLogs.get(topicPartition) - if (log != null) { - return log + if (!snapshotRead && elasticLogs.containsKey(topicPartition)) { + return elasticLogs.get(topicPartition) } var elasticLog: ElasticUnifiedLog = null // Only Partition#makeLeader will create a new log, the ReplicaManager#asyncApplyDelta will ensure the same partition @@ -71,9 +79,12 @@ class ElasticLogManager(val client: Client, val openStreamChecker: OpenStreamChe logOffsetsListener = LogOffsetsListener.NO_OP_OFFSETS_LISTENER, client, NAMESPACE, - openStreamChecker + openStreamChecker, + OpenHint.isSnapshotRead ) - elasticLogs.put(topicPartition, elasticLog) + if (!snapshotRead) { + elasticLogs.put(topicPartition, elasticLog) + } elasticLog } @@ -96,11 +107,9 @@ class ElasticLogManager(val client: Client, val openStreamChecker: OpenStreamChe /** * Remove elastic log in the map. - * - * @param topicPartition topic partition */ - def removeLog(topicPartition: TopicPartition): Unit = { - elasticLogs.remove(topicPartition) + def removeLog(topicPartition: TopicPartition, log: ElasticUnifiedLog): Unit = { + elasticLogs.remove(topicPartition, log) } def startup(): Unit = { @@ -132,7 +141,7 @@ object ElasticLogManager { context.config = config context.brokerServer = broker val openStreamChecker = if (broker != null) { - new DefaultOpenStreamChecker(broker.metadataCache) + new DefaultOpenStreamChecker(config.nodeId, broker.metadataCache) } else { OpenStreamChecker.NOOP } @@ -154,8 +163,8 @@ object ElasticLogManager { def enabled(): Boolean = isEnabled - def removeLog(topicPartition: TopicPartition): Unit = { - instance().get.removeLog(topicPartition) + def removeLog(topicPartition: TopicPartition, log: ElasticUnifiedLog): Unit = { + instance().get.removeLog(topicPartition, log) } def destroyLog(topicPartition: TopicPartition, topicId: Uuid, epoch: Long): Unit = { @@ -173,7 +182,7 @@ object ElasticLogManager { } // visible for testing - def getOrCreateLog(dir: File, + def createLog(dir: File, config: LogConfig, scheduler: Scheduler, time: Time, @@ -184,7 +193,7 @@ object ElasticLogManager { logDirFailureChannel: LogDirFailureChannel, topicId: Option[Uuid], leaderEpoch: Long = 0): ElasticUnifiedLog = { - instance().get.getOrCreateLog( + instance().get.createLog( dir, config, scheduler, diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogMeta.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogMeta.java index 34dc84df1c..caf2e018f1 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogMeta.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogMeta.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegment.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegment.java index 543d31d905..54e5094dca 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegment.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegment.java @@ -1,16 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; +import kafka.cluster.PartitionSnapshot; import kafka.log.streamaspect.cache.FileCache; import org.apache.kafka.common.InvalidRecordException; @@ -31,6 +40,7 @@ import org.apache.kafka.storage.internals.log.LogFileUtils; import org.apache.kafka.storage.internals.log.LogOffsetMetadata; import org.apache.kafka.storage.internals.log.LogSegment; +import org.apache.kafka.storage.internals.log.LogSegmentOffsetOverflowException; import org.apache.kafka.storage.internals.log.OffsetIndex; import org.apache.kafka.storage.internals.log.OffsetPosition; import org.apache.kafka.storage.internals.log.ProducerAppendInfo; @@ -186,9 +196,26 @@ public int size() { return log.sizeInBytes(); } + /** + * Checks that the argument offset can be represented as an integer offset relative to the baseOffset. + * This method is similar in purpose to {@see org.apache.kafka.storage.internals.log.LogSegment#canConvertToRelativeOffset}. + *

    + * The implementation is inspired by {@see org.apache.kafka.storage.internals.log.AbstractIndex#canAppendOffset}, + * but uses {@code < Integer.MAX_VALUE} instead of {@code <= Integer.MAX_VALUE} to address an offset overflow issue. + * + * @param offset The offset to check. + * @return true if the offset can be converted, false otherwise. + * @see Issue #2718 + */ private boolean canConvertToRelativeOffset(long offset) { long relativeOffset = offset - baseOffset; - return relativeOffset >= 0 && relativeOffset <= Integer.MAX_VALUE; + // Note: The check is `relativeOffset < Integer.MAX_VALUE` instead of `<=` to avoid overflow. + // See https://github.com/AutoMQ/automq/issues/2718 for details. + return relativeOffset >= 0 && relativeOffset < Integer.MAX_VALUE; + } + private void ensureOffsetInRange(long offset) throws IOException { + if (!canConvertToRelativeOffset(offset)) + throw new LogSegmentOffsetOverflowException(this, offset); } @Override @@ -208,9 +235,13 @@ public void append( meta.firstBatchTimestamp(largestTimestampMs); } + ensureOffsetInRange(largestOffset); + // append the messages long appendedBytes = log.append(records, largestOffset + 1); - LOGGER.trace("Appended {} to {} at end offset {}", appendedBytes, log, largestOffset); + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("Appended {} to {} at end offset {}", appendedBytes, log, largestOffset); + } // Update the in memory max timestamp and corresponding offset. if (largestTimestampMs > maxTimestampSoFar()) { maxTimestampAndOffsetSoFar = new TimestampOffset(largestTimestampMs, offsetOfMaxTimestamp); @@ -262,6 +293,10 @@ public CompletableFuture readAsync(long startOffset, int maxSize, return CompletableFuture.failedFuture(new IllegalArgumentException("Invalid max size " + maxSize + " for log read from segment " + log)); // Note that relativePositionInSegment here is a fake value. There are no 'position' in elastic streams. + // if the start offset is less than base offset, use base offset. This usually happens when the prev segment is generated + // by compaction and its last offset is less than the base offset of the current segment. + startOffset = Utils.max(startOffset, baseOffset); + // if the start offset is already off the end of the log, return null if (startOffset >= log.nextOffset()) { return CompletableFuture.completedFuture(null); @@ -275,7 +310,7 @@ public CompletableFuture readAsync(long startOffset, int maxSize, return CompletableFuture.completedFuture(new FetchDataInfo(offsetMetadata, MemoryRecords.EMPTY)); } - return log.read(startOffset, maxOffset, maxSize) + return log.read(startOffset, maxOffset, adjustedMaxSize) .thenApply(records -> new FetchDataInfo(offsetMetadata, records)); } @@ -575,4 +610,9 @@ public Record next() { }; }; } + + void snapshot(PartitionSnapshot snapshot) { + log.size(snapshot.logEndOffset().relativePositionInSegment); + timeIndex.snapshot(snapshot.lastTimestampOffset()); + } } diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEvent.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEvent.java index 6f2405a4f2..1fd4fcb62f 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEvent.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEvent.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEventListener.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEventListener.java index 11b6c77872..7ac6886f15 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEventListener.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentEventListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentManager.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentManager.java index e92c1ba40e..d382f8e1e8 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentManager.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogSegmentManager.java @@ -1,25 +1,35 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; -import org.apache.kafka.storage.internals.log.LogOffsetMetadata; +import kafka.cluster.LogEventListener; + +import org.apache.kafka.storage.internals.log.LogSegment; import com.automq.stream.api.Stream; -import com.automq.stream.utils.Threads; import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collection; import java.util.HashMap; import java.util.Iterator; import java.util.List; @@ -28,7 +38,7 @@ import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; @@ -40,13 +50,15 @@ public class ElasticLogSegmentManager { private final ReentrantLock segmentLock = new ReentrantLock(); private final Map segments = new HashMap<>(); private final Map inflightCleanedSegments = new HashMap<>(); - private final EventListener segmentEventListener = new EventListener(); - final AtomicReference offsetUpperBound = new AtomicReference<>(); + private final EventListener innerListener = new EventListener(); + private final List logEventListeners = new CopyOnWriteArrayList<>(); private final MetaStream metaStream; private final ElasticLogStreamManager streamManager; private final String logIdent; + private volatile ElasticLogMeta logMeta; + public ElasticLogSegmentManager(MetaStream metaStream, ElasticLogStreamManager streamManager, String logIdent) { this.metaStream = metaStream; this.streamManager = streamManager; @@ -61,6 +73,7 @@ public void put(long baseOffset, ElasticLogSegment segment) { } finally { segmentLock.unlock(); } + notifyLogEventListeners(segment, LogEventListener.Event.SEGMENT_CREATE); } public void putInflightCleaned(long baseOffset, ElasticLogSegment segment) { @@ -73,26 +86,25 @@ public void putInflightCleaned(long baseOffset, ElasticLogSegment segment) { } public CompletableFuture create(long baseOffset, ElasticLogSegment segment) { - LogOffsetMetadata offset = new LogOffsetMetadata(baseOffset, baseOffset, 0); - while (!offsetUpperBound.compareAndSet(null, offset)) { - LOGGER.info("{} try create new segment with offset $baseOffset, wait last segment meta persisted.", logIdent); - Threads.sleep(1L); - } segmentLock.lock(); try { segments.put(baseOffset, segment); } finally { segmentLock.unlock(); } - return asyncPersistLogMeta().thenAccept(nil -> { - offsetUpperBound.set(null); - }); + return asyncPersistLogMeta().thenAccept(rst -> + notifyLogEventListeners(segment, LogEventListener.Event.SEGMENT_CREATE) + ); } public ElasticLogSegment remove(long baseOffset) { segmentLock.lock(); try { - return segments.remove(baseOffset); + ElasticLogSegment segment = segments.remove(baseOffset); + if (segment != null) { + notifyLogEventListeners(segment, LogEventListener.Event.SEGMENT_DELETE); + } + return segment; } finally { segmentLock.unlock(); } @@ -119,6 +131,7 @@ public CompletableFuture asyncPersistLogMeta() { .collect(Collectors.toList()); meta = logMeta(streams, segmentList); + this.logMeta = meta; // We calculate trimOffsets in the lock to ensure that no more new stream with data is created during the calculation. trimOffsets = calTrimOffset( streams, @@ -186,7 +199,53 @@ private static void calTrimOffset(Map streamMinOffsets, ElasticStr } public ElasticLogSegmentEventListener logSegmentEventListener() { - return segmentEventListener; + return innerListener; + } + + public ElasticLogMeta logMeta() { + if (logMeta == null) { + segmentLock.lock(); + try { + if (logMeta != null) { + return logMeta; + } + logMeta = generateLogMeta(); + } finally { + segmentLock.unlock(); + } + } + return logMeta; + } + + public Collection streams() { + return streamManager.streams().values(); + } + + public void addLogEventListener(LogEventListener listener) { + logEventListeners.add(listener); + } + + public void notifySegmentUpdate() { + notifyLogEventListeners(null, LogEventListener.Event.SEGMENT_UPDATE); + } + + private void notifyLogEventListeners(LogSegment segment, LogEventListener.Event event) { + for (LogEventListener listener : logEventListeners) { + try { + listener.onChanged(segment, event); + } catch (Throwable e) { + LOGGER.error("got notify listener error", e); + } + } + } + + private ElasticLogMeta generateLogMeta() { + Map streams = streamManager.streams(); + List segmentList = segments.values().stream() + .sorted() + .map(ElasticLogSegment::meta) + .collect(Collectors.toList()); + return logMeta(streams, segmentList); } public static ElasticLogMeta logMeta(Map streams, List segmentList) { diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticLogStreamManager.java b/core/src/main/scala/kafka/log/streamaspect/ElasticLogStreamManager.java index 1a3ec67cf0..c44fb8cbce 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticLogStreamManager.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticLogStreamManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -30,6 +38,7 @@ public class ElasticLogStreamManager { private final int replicaCount; private final long epoch; private final Map tags; + private final boolean snapshotRead; /** * inner listener for created LazyStream */ @@ -39,15 +48,16 @@ public class ElasticLogStreamManager { */ private ElasticStreamEventListener outerListener; - public ElasticLogStreamManager(Map streams, StreamClient streamClient, int replicaCount, long epoch, Map tags) throws IOException { + public ElasticLogStreamManager(Map streams, StreamClient streamClient, int replicaCount, long epoch, Map tags, boolean snapshotRead) throws IOException { this.streamClient = streamClient; this.replicaCount = replicaCount; this.epoch = epoch; this.tags = tags; + this.snapshotRead = snapshotRead; for (Map.Entry entry : streams.entrySet()) { String name = entry.getKey(); long streamId = entry.getValue(); - LazyStream stream = new LazyStream(name, streamId, streamClient, replicaCount, epoch, tags); + LazyStream stream = new LazyStream(name, streamId, streamClient, replicaCount, epoch, tags, snapshotRead); stream.setListener(innerListener); streamMap.put(name, stream); } @@ -57,7 +67,10 @@ public LazyStream getStream(String name) throws IOException { if (streamMap.containsKey(name)) { return streamMap.get(name); } - LazyStream lazyStream = new LazyStream(name, LazyStream.NOOP_STREAM_ID, streamClient, replicaCount, epoch, tags); + if (snapshotRead) { + throw new IllegalStateException("snapshotRead mode can not create stream"); + } + LazyStream lazyStream = new LazyStream(name, LazyStream.NOOP_STREAM_ID, streamClient, replicaCount, epoch, tags, snapshotRead); lazyStream.setListener(innerListener); // pre-create log and tim stream cause of their high frequency of use. boolean warmUp = "log".equals(name) || "tim".equals(name); @@ -72,6 +85,24 @@ public Map streams() { return Collections.unmodifiableMap(streamMap); } + public void createIfNotExist(String name, Long streamId) { + LazyStream s = streamMap.get(name); + if (s != null && s.streamId() == streamId) { + // precheck to prevent create lambda. + return; + } + streamMap.compute(name, (n, stream) -> { + if (stream != null && stream.streamId() == streamId) { + return stream; + } + try { + return new LazyStream(name, streamId, streamClient, replicaCount, epoch, tags, snapshotRead); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } + public void setListener(ElasticStreamEventListener listener) { this.outerListener = listener; } diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionMeta.java b/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionMeta.java index 7a620a7522..96e7a1f268 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionMeta.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionMeta.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionProducerSnapshotsMeta.java b/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionProducerSnapshotsMeta.java index 3b577ebf8e..0806dfcaad 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionProducerSnapshotsMeta.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticPartitionProducerSnapshotsMeta.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticProducerStateManager.java b/core/src/main/scala/kafka/log/streamaspect/ElasticProducerStateManager.java index 56ef4e6ac5..e37004653a 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticProducerStateManager.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticProducerStateManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticResourceStatus.java b/core/src/main/scala/kafka/log/streamaspect/ElasticResourceStatus.java index 44d027fcb2..f713dee7ac 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticResourceStatus.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticResourceStatus.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamEventListener.java b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamEventListener.java index 685dbfb7f7..76f16dc9db 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamEventListener.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamEventListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamMetaEvent.java b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamMetaEvent.java index 1d8aa1c86d..2052485c4b 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamMetaEvent.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamMetaEvent.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSegmentMeta.java b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSegmentMeta.java index e5eda4f515..fd41584f64 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSegmentMeta.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSegmentMeta.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSlice.java b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSlice.java index fe6ee7fccd..bb27d5ecc2 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSlice.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSlice.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -66,12 +74,6 @@ default CompletableFuture fetch(long startOffset, long endOffset) { */ long confirmOffset(); - /** - * Get slice start offset in under stream. - * - * @return segment start offset in stream. - */ - long startOffsetInStream(); /** * Get slice range which is the relative offset range in stream. diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSliceManager.java b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSliceManager.java index c121795286..8e5ff90b82 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSliceManager.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticStreamSliceManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticTimeIndex.java b/core/src/main/scala/kafka/log/streamaspect/ElasticTimeIndex.java index b3521d4eac..10b7aa0a4b 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticTimeIndex.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticTimeIndex.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -312,4 +320,13 @@ private TimestampOffset lastEntryFromIndexFile() { return entry(entries() - 1); } } + + void snapshot(TimestampOffset lastTimestampOffset) { + if (lastTimestampOffset == null) { + return; + } + setEntries((int) (stream.nextOffset() / ENTRY_SIZE)); + lastEntry(lastTimestampOffset); + } + } diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticTransactionIndex.java b/core/src/main/scala/kafka/log/streamaspect/ElasticTransactionIndex.java index 654f3dca5a..67f1a24b87 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticTransactionIndex.java +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticTransactionIndex.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -54,7 +62,7 @@ public ElasticTransactionIndex(long startOffset, File file, StreamSliceSupplier this.cache = cache; this.cacheId = cache.newCacheId(); this.path = file.getPath(); - lastAppend = new LastAppend(stream.nextOffset(), CompletableFuture.completedFuture(null)); + lastAppend = new LastAppend(CompletableFuture.completedFuture(null)); } @Override @@ -75,7 +83,7 @@ public void append(AbortedTxn abortedTxn) { lastOffset = OptionalLong.of(abortedTxn.lastOffset()); long position = stream.nextOffset(); CompletableFuture cf = stream.append(RawPayloadRecordBatch.of(abortedTxn.buffer().duplicate())); - lastAppend = new LastAppend(stream.nextOffset(), cf); + lastAppend = new LastAppend(cf); cache.put(cacheId, position, Unpooled.wrappedBuffer(abortedTxn.buffer())); } @@ -160,12 +168,12 @@ public void seal() { protected Iterable iterable(Supplier allocate) { // await last append complete, usually the abort transaction is not frequent, so it's ok to block here. LastAppend lastAppend = this.lastAppend; - int endPosition = (int) lastAppend.offset; try { lastAppend.cf.get(); } catch (Throwable e) { throw new RuntimeException(e); } + int endPosition = (int) stream.confirmOffset(); PrimitiveRef.IntRef position = PrimitiveRef.ofInt(0); Queue queue = new ArrayDeque<>(); return () -> new Iterator<>() { @@ -218,11 +226,9 @@ private FetchResult fetchStream(long startOffset, long endOffset, int maxBytes) } static class LastAppend { - final long offset; final CompletableFuture cf; - LastAppend(long offset, CompletableFuture cf) { - this.offset = offset; + LastAppend(CompletableFuture cf) { this.cf = cf; } diff --git a/core/src/main/scala/kafka/log/streamaspect/ElasticUnifiedLog.scala b/core/src/main/scala/kafka/log/streamaspect/ElasticUnifiedLog.scala index 44d425bc13..17639ca5be 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ElasticUnifiedLog.scala +++ b/core/src/main/scala/kafka/log/streamaspect/ElasticUnifiedLog.scala @@ -1,18 +1,27 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect import com.automq.stream.api.Client -import com.automq.stream.utils.FutureUtil +import com.automq.stream.utils.{FutureUtil, Threads} +import kafka.cluster.PartitionSnapshot import kafka.log._ import kafka.log.streamaspect.ElasticUnifiedLog.{CheckpointExecutor, MaxCheckpointIntervalBytes, MinCheckpointIntervalMs} import kafka.server._ @@ -20,9 +29,9 @@ import kafka.utils.Logging import org.apache.kafka.common.errors.OffsetOutOfRangeException import org.apache.kafka.common.errors.s3.StreamFencedException import org.apache.kafka.common.record.{MemoryRecords, RecordVersion} -import org.apache.kafka.common.utils.{ThreadUtils, Time} +import org.apache.kafka.common.utils.{Time, Utils} import org.apache.kafka.common.{TopicPartition, Uuid} -import org.apache.kafka.server.common.MetadataVersion +import org.apache.kafka.server.common.{MetadataVersion, OffsetAndEpoch} import org.apache.kafka.server.util.Scheduler import org.apache.kafka.storage.internals.epoch.LeaderEpochFileCache import org.apache.kafka.storage.internals.log._ @@ -32,7 +41,7 @@ import java.nio.ByteBuffer import java.nio.file.Path import java.util import java.util.concurrent.atomic.LongAdder -import java.util.concurrent.{CompletableFuture, ConcurrentHashMap, Executors} +import java.util.concurrent.{CompletableFuture, ConcurrentHashMap, CopyOnWriteArrayList} import scala.jdk.CollectionConverters.CollectionHasAsScala import scala.util.{Failure, Success, Try} @@ -43,13 +52,12 @@ class ElasticUnifiedLog(_logStartOffset: Long, _leaderEpochCache: Option[LeaderEpochFileCache], producerStateManager: ProducerStateManager, __topicId: Option[Uuid], - logOffsetsListener: LogOffsetsListener + logOffsetsListener: LogOffsetsListener, + var snapshotRead: Boolean, ) extends UnifiedLog(_logStartOffset, elasticLog, brokerTopicStats, producerIdExpirationCheckIntervalMs, _leaderEpochCache, producerStateManager, __topicId, false, false, logOffsetsListener) { - ElasticUnifiedLog.Logs.put(elasticLog.topicPartition, this) - var confirmOffsetChangeListener: Option[() => Unit] = None elasticLog.confirmOffsetChangeListener = Some(() => confirmOffsetChangeListener.map(_.apply())) @@ -57,6 +65,13 @@ class ElasticUnifiedLog(_logStartOffset: Long, // fuzzy interval bytes for checkpoint, it's ok not thread safe var checkpointIntervalBytes = 0 var lastCheckpointTimestamp = time.milliseconds() + var configChangeListeners = new CopyOnWriteArrayList[LogConfigChangeListener]() + + override def newMetrics(): Unit = { + if (!snapshotRead) { + super.newMetrics() + } + } def getLocalLog(): ElasticLog = elasticLog @@ -84,11 +99,14 @@ class ElasticUnifiedLog(_logStartOffset: Long, rst } - def tryCheckpoint(): Unit = { + def tryCheckpoint(): Boolean = { if (checkpointIntervalBytes > 0) { checkpointIntervalBytes = 0 lastCheckpointTimestamp = time.milliseconds() checkpoint() + true + } else { + false } } @@ -197,11 +215,12 @@ class ElasticUnifiedLog(_logStartOffset: Long, flush(true) elasticLog.close() } - elasticLog.segments.clear() // graceful await append ack elasticLog.lastAppendAckFuture.get() elasticLog.isMemoryMappedBufferClosed = true - elasticLog.deleteEmptyDir() + // Since https://github.com/AutoMQ/automq/pull/2837 , AutoMQ won't create the partition directory when the partition opens + // The deletion here aims to clean the old directory. + Utils.delete(dir) } /** @@ -226,15 +245,74 @@ class ElasticUnifiedLog(_logStartOffset: Long, // noop implementation, producer snapshot and recover point will be appended to MetaStream, so they have order relation. } + override def updateConfig( + newConfig: LogConfig): LogConfig = { + val config = super.updateConfig(newConfig) + for (listener <- configChangeListeners.asScala) { + try { + listener.onLogConfigChange(this, newConfig) + } catch { + case e: Throwable => + error(s"Error while invoking config change listener $listener", e) + } + } + config + } + + override def endOffsetForEpoch(leaderEpoch: Int): Option[OffsetAndEpoch] = { + if (snapshotRead) { + Option(new OffsetAndEpoch(logEndOffset, leaderEpoch)) + } else { + super.endOffsetForEpoch(leaderEpoch) + } + } + // only used for test def listProducerSnapshots(): util.NavigableMap[java.lang.Long, ByteBuffer] = { producerStateManager.asInstanceOf[ElasticProducerStateManager].snapshotsMap } + + def addConfigChangeListener(listener: LogConfigChangeListener): Unit = { + configChangeListeners.add(listener) + } + + def snapshot(snapshot: PartitionSnapshot.Builder): Unit = { + lock synchronized { + snapshot.firstUnstableOffset(firstUnstableOffsetMetadata.orNull) + val localLog = getLocalLog() + localLog.snapshot(snapshot) + } + } + + def snapshot(snapshot: PartitionSnapshot): Unit = { + lock synchronized { + val localLog = getLocalLog() + localLog.snapshot(snapshot) + if (snapshot.firstUnstableOffset() == null) { + firstUnstableOffsetMetadata = None + } else { + var offset = snapshot.firstUnstableOffset() + val segmentBaseOffset = localLog.segments.floorSegment(offset.messageOffset).get().baseOffset() + offset = new LogOffsetMetadata(offset.messageOffset, segmentBaseOffset, offset.relativePositionInSegment) + firstUnstableOffsetMetadata = Some(offset) + } + if (snapshot.logMeta() != null) { + val opt = localLog.segments.firstSegmentBaseOffset() + opt.ifPresent(baseOffset => { + updateLogStartOffset(baseOffset) + }) + } + highWatermarkMetadata = localLog.logEndOffsetMetadata + } + } + } object ElasticUnifiedLog extends Logging { - private val CheckpointExecutor = Executors.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("checkpoint-executor", true)) - private val MaxCheckpointIntervalBytes = 50 * 1024 * 1024 + private val CheckpointExecutor = { + Threads.newSingleThreadScheduledExecutor("checkpoint-executor", true, logger.underlying) + } + private val MaxCheckpointIntervalBytes = 50 * 1024 * 1024 private val MinCheckpointIntervalMs = 10 * 1000 private val Logs = new ConcurrentHashMap[TopicPartition, ElasticUnifiedLog]() // fuzzy dirty bytes for checkpoint, it's ok not thread safe @@ -256,10 +334,10 @@ object ElasticUnifiedLog extends Logging { topicId: Option[Uuid], leaderEpoch: Long = 0, logOffsetsListener: LogOffsetsListener, - client: Client, namespace: String, openStreamChecker: OpenStreamChecker, + snapshotRead: Boolean = false ): ElasticUnifiedLog = { val topicPartition = UnifiedLog.parseTopicPartitionName(dir) val partitionLogDirFailureChannel = new PartitionLogDirFailureChannel(logDirFailureChannel, dir.getPath); @@ -268,8 +346,9 @@ object ElasticUnifiedLog extends Logging { var localLog: ElasticLog = null while(localLog == null) { try { - localLog = ElasticLog(client, namespace, dir, config, scheduler, time, topicPartition, partitionLogDirFailureChannel, - new ConcurrentHashMap[String, Int](), maxTransactionTimeoutMs, producerStateManagerConfig, topicId, leaderEpoch, openStreamChecker) + localLog = ElasticLog(client, namespace, dir, config, scheduler, time, topicPartition, + partitionLogDirFailureChannel, new ConcurrentHashMap[String, Int](), maxTransactionTimeoutMs, + producerStateManagerConfig, topicId, leaderEpoch, openStreamChecker, snapshotRead) } catch { case e: Throwable => val cause = FutureUtil.cause(e) @@ -291,10 +370,16 @@ object ElasticUnifiedLog extends Logging { _leaderEpochCache = leaderEpochFileCache, localLog.producerStateManager, topicId, - logOffsetsListener + logOffsetsListener, + snapshotRead ) val timeCost = System.currentTimeMillis() - start info(s"ElasticUnifiedLog $topicPartition opened time cost: $timeCost ms") + + if (!snapshotRead) { + ElasticUnifiedLog.Logs.put(elasticUnifiedLog.getLocalLog().topicPartition, elasticUnifiedLog) + } + elasticUnifiedLog } } @@ -306,7 +391,10 @@ object ElasticUnifiedLog extends Logging { DirtyBytes.reset() for (log <- Logs.values().asScala) { try { - log.tryCheckpoint() + if (log.tryCheckpoint()) { + // sleep a while to avoid too many checkpoint at the same time, which may cause high append latency + Thread.sleep(10) + } } catch { case e: Throwable => error("Error while checkpoint", e) } diff --git a/core/src/main/scala/kafka/log/streamaspect/LazyStream.java b/core/src/main/scala/kafka/log/streamaspect/LazyStream.java index 67d255c762..e8ddb49885 100644 --- a/core/src/main/scala/kafka/log/streamaspect/LazyStream.java +++ b/core/src/main/scala/kafka/log/streamaspect/LazyStream.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -47,7 +55,7 @@ public class LazyStream implements Stream { private volatile Stream inner = NOOP_STREAM; private ElasticStreamEventListener eventListener; - public LazyStream(String name, long streamId, StreamClient client, int replicaCount, long epoch, Map tags) throws IOException { + public LazyStream(String name, long streamId, StreamClient client, int replicaCount, long epoch, Map tags, boolean snapshotRead) throws IOException { this.name = name; this.client = client; this.replicaCount = replicaCount; @@ -56,7 +64,11 @@ public LazyStream(String name, long streamId, StreamClient client, int replicaCo if (streamId != NOOP_STREAM_ID) { try { // open exist stream - inner = client.openStream(streamId, OpenStreamOptions.builder().epoch(epoch).tags(tags).build()).get(); + OpenStreamOptions.Builder options = OpenStreamOptions.builder().epoch(epoch).tags(tags); + if (snapshotRead) { + options.readWriteMode(OpenStreamOptions.ReadWriteMode.SNAPSHOT_READ); + } + inner = client.openStream(streamId, options.build()).get(); } catch (InterruptedException e) { throw new RuntimeException(e); } catch (ExecutionException e) { @@ -102,6 +114,11 @@ public long confirmOffset() { return inner.confirmOffset(); } + @Override + public void confirmOffset(long offset) { + inner.confirmOffset(offset); + } + @Override public long nextOffset() { return inner.nextOffset(); @@ -142,6 +159,11 @@ public CompletableFuture destroy() { return inner.destroy(); } + @Override + public CompletableFuture lastAppendFuture() { + return inner.lastAppendFuture(); + } + @Override public String toString() { return "LazyStream{" + "name='" + name + '\'' + "streamId='" + inner.streamId() + '\'' + ", replicaCount=" + replicaCount + '}'; @@ -187,6 +209,10 @@ public long confirmOffset() { return 0; } + @Override + public void confirmOffset(long offset) { + } + @Override public long nextOffset() { return 0; @@ -216,5 +242,10 @@ public CompletableFuture close() { public CompletableFuture destroy() { return CompletableFuture.completedFuture(null); } + + @Override + public CompletableFuture lastAppendFuture() { + return null; + } } } diff --git a/core/src/main/scala/kafka/log/streamaspect/LogConfigChangeListener.java b/core/src/main/scala/kafka/log/streamaspect/LogConfigChangeListener.java new file mode 100644 index 0000000000..24e3a652ef --- /dev/null +++ b/core/src/main/scala/kafka/log/streamaspect/LogConfigChangeListener.java @@ -0,0 +1,28 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.streamaspect; + +import org.apache.kafka.storage.internals.log.LogConfig; + +public interface LogConfigChangeListener { + + void onLogConfigChange(ElasticUnifiedLog log, LogConfig logConfig); + +} diff --git a/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java b/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java index 52fc068dd9..cae0cc6249 100644 --- a/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java +++ b/core/src/main/scala/kafka/log/streamaspect/MemoryClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -72,10 +80,11 @@ public CompletableFuture failover(FailoverRequest request) { return FutureUtil.failedFuture(new UnsupportedOperationException()); } - static class StreamImpl implements Stream { + public static class StreamImpl implements Stream { private final AtomicLong nextOffsetAlloc = new AtomicLong(); private NavigableMap recordMap = new ConcurrentSkipListMap<>(); private final long streamId; + private volatile CompletableFuture lastAppendFuture; public StreamImpl(long streamId) { this.streamId = streamId; @@ -101,6 +110,11 @@ public long confirmOffset() { return nextOffsetAlloc.get(); } + @Override + public void confirmOffset(long offset) { + nextOffsetAlloc.set(offset); + } + @Override public long nextOffset() { return nextOffsetAlloc.get(); @@ -114,7 +128,8 @@ public synchronized CompletableFuture append(AppendContext context copy.flip(); recordBatch = new DefaultRecordBatch(recordBatch.count(), recordBatch.baseTimestamp(), recordBatch.properties(), copy); recordMap.put(baseOffset, new RecordBatchWithContextWrapper(recordBatch, baseOffset)); - return CompletableFuture.completedFuture(() -> baseOffset); + this.lastAppendFuture = CompletableFuture.completedFuture(() -> baseOffset); + return lastAppendFuture; } @Override @@ -124,10 +139,26 @@ public CompletableFuture fetch(FetchContext context, long startOffs if (floorKey == null) { return CompletableFuture.completedFuture(ArrayList::new); } - List records = new ArrayList<>(recordMap.subMap(floorKey, endOffset).values()); + NavigableMap subMap = recordMap.subMap(floorKey, true, endOffset, false); + List records = new ArrayList<>(); + int accumulatedSize = 0; + for (Map.Entry entry : subMap.entrySet()) { + RecordBatchWithContext batch = entry.getValue(); + int batchSize = batch.rawPayload().remaining(); + if (accumulatedSize + batchSize > maxSizeHint && !records.isEmpty()) { + break; + } + records.add(batch); + accumulatedSize += batchSize; + + if (accumulatedSize > maxSizeHint) { + break; + } + } return CompletableFuture.completedFuture(() -> records); } + @Override public CompletableFuture trim(long newStartOffset) { recordMap = new ConcurrentSkipListMap<>(recordMap.tailMap(newStartOffset)); @@ -144,6 +175,11 @@ public CompletableFuture destroy() { recordMap.clear(); return CompletableFuture.completedFuture(null); } + + @Override + public CompletableFuture lastAppendFuture() { + return lastAppendFuture; + } } static class StreamClientImpl implements StreamClient { diff --git a/core/src/main/scala/kafka/log/streamaspect/MetaKeyValue.java b/core/src/main/scala/kafka/log/streamaspect/MetaKeyValue.java index 7781f6c503..d078fab3ab 100644 --- a/core/src/main/scala/kafka/log/streamaspect/MetaKeyValue.java +++ b/core/src/main/scala/kafka/log/streamaspect/MetaKeyValue.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/MetaStream.java b/core/src/main/scala/kafka/log/streamaspect/MetaStream.java index 1d1acc21e7..0818ab04cb 100644 --- a/core/src/main/scala/kafka/log/streamaspect/MetaStream.java +++ b/core/src/main/scala/kafka/log/streamaspect/MetaStream.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -31,6 +39,7 @@ import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.OptionalLong; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; @@ -108,6 +117,11 @@ public long confirmOffset() { return innerStream.confirmOffset(); } + @Override + public void confirmOffset(long offset) { + innerStream.confirmOffset(offset); + } + @Override public long nextOffset() { return innerStream.nextOffset(); @@ -182,6 +196,11 @@ public CompletableFuture destroy() { return innerStream.destroy(); } + @Override + public CompletableFuture lastAppendFuture() { + return innerStream.lastAppendFuture(); + } + /** * Replay meta stream and return a map of meta keyValues. KeyValues will be cached in metaCache. * @@ -244,6 +263,10 @@ public Map replay() throws IOException { return getValidMetaMap(); } + public Optional get(String key) { + return Optional.ofNullable(metaCache.get(key)).map(o -> o.value.slice()); + } + private Map getValidMetaMap() { Map metaMap = new HashMap<>(); metaCache.forEach((key, value) -> { @@ -261,7 +284,7 @@ private Map getValidMetaMap() { metaMap.put(key, ElasticLeaderEpochCheckpointMeta.decode(value.value())); break; default: - LOGGER.error("{} streamId {}: unknown meta key: {}", logIdent, streamId(), key); + metaMap.put(key, value.value().duplicate()); } }); return metaMap; @@ -311,10 +334,9 @@ private synchronized CompletableFuture doCompaction(boolean force) { return CompletableFuture.completedFuture(null); } CompletableFuture overwriteCf = CompletableFuture.allOf(overwrite.stream().map(this::append).toArray(CompletableFuture[]::new)); - return overwriteCf.thenAccept(nil -> { - OptionalLong minOffset = metaCache.values().stream().mapToLong(v -> v.offset).min(); - minOffset.ifPresent(this::trim); - }); + OptionalLong minOffset = metaCache.values().stream().mapToLong(v -> v.offset).min(); + // await overwrite complete then trim to the minimum offset in metaCache + return overwriteCf.thenAccept(nil -> minOffset.ifPresent(this::trim)); } static class MetadataValue { diff --git a/core/src/main/scala/kafka/log/streamaspect/Offsets.java b/core/src/main/scala/kafka/log/streamaspect/Offsets.java index 4e3a61bce5..6938defed7 100644 --- a/core/src/main/scala/kafka/log/streamaspect/Offsets.java +++ b/core/src/main/scala/kafka/log/streamaspect/Offsets.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/OpenHint.java b/core/src/main/scala/kafka/log/streamaspect/OpenHint.java new file mode 100644 index 0000000000..9635a03571 --- /dev/null +++ b/core/src/main/scala/kafka/log/streamaspect/OpenHint.java @@ -0,0 +1,43 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.streamaspect; + +import io.netty.util.concurrent.FastThreadLocal; + +public class OpenHint { + public static final FastThreadLocal SNAPSHOT_READ = new FastThreadLocal<>() { + @Override + protected Boolean initialValue() { + return false; + } + }; + + public static void markSnapshotRead() { + SNAPSHOT_READ.set(true); + } + + public static boolean isSnapshotRead() { + return SNAPSHOT_READ.get(); + } + + public static void clear() { + SNAPSHOT_READ.remove(); + } +} diff --git a/core/src/main/scala/kafka/log/streamaspect/OpenStreamChecker.java b/core/src/main/scala/kafka/log/streamaspect/OpenStreamChecker.java index 0ad03465e0..e0b6757023 100644 --- a/core/src/main/scala/kafka/log/streamaspect/OpenStreamChecker.java +++ b/core/src/main/scala/kafka/log/streamaspect/OpenStreamChecker.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/PartitionLogDirFailureChannel.java b/core/src/main/scala/kafka/log/streamaspect/PartitionLogDirFailureChannel.java index 28cf361268..8c0c4fbf95 100644 --- a/core/src/main/scala/kafka/log/streamaspect/PartitionLogDirFailureChannel.java +++ b/core/src/main/scala/kafka/log/streamaspect/PartitionLogDirFailureChannel.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/PartitionStatusTracker.java b/core/src/main/scala/kafka/log/streamaspect/PartitionStatusTracker.java index 15d9f39139..bbc78e6f4d 100644 --- a/core/src/main/scala/kafka/log/streamaspect/PartitionStatusTracker.java +++ b/core/src/main/scala/kafka/log/streamaspect/PartitionStatusTracker.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/RawKafkaMeta/RawKafkaMeta.scala b/core/src/main/scala/kafka/log/streamaspect/RawKafkaMeta/RawKafkaMeta.scala index 75587dec9b..30cc72a565 100644 --- a/core/src/main/scala/kafka/log/streamaspect/RawKafkaMeta/RawKafkaMeta.scala +++ b/core/src/main/scala/kafka/log/streamaspect/RawKafkaMeta/RawKafkaMeta.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect diff --git a/core/src/main/scala/kafka/log/streamaspect/RawPayloadRecordBatch.java b/core/src/main/scala/kafka/log/streamaspect/RawPayloadRecordBatch.java index 4ba1f11271..aa2570f26e 100644 --- a/core/src/main/scala/kafka/log/streamaspect/RawPayloadRecordBatch.java +++ b/core/src/main/scala/kafka/log/streamaspect/RawPayloadRecordBatch.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/ReadHint.java b/core/src/main/scala/kafka/log/streamaspect/ReadHint.java index dd00ee157b..08747b106f 100644 --- a/core/src/main/scala/kafka/log/streamaspect/ReadHint.java +++ b/core/src/main/scala/kafka/log/streamaspect/ReadHint.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/SliceRange.java b/core/src/main/scala/kafka/log/streamaspect/SliceRange.java index 406013e5ff..8ab1d94543 100644 --- a/core/src/main/scala/kafka/log/streamaspect/SliceRange.java +++ b/core/src/main/scala/kafka/log/streamaspect/SliceRange.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/StreamSliceSupplier.java b/core/src/main/scala/kafka/log/streamaspect/StreamSliceSupplier.java index 0639d6f0a4..a8e4adbc16 100644 --- a/core/src/main/scala/kafka/log/streamaspect/StreamSliceSupplier.java +++ b/core/src/main/scala/kafka/log/streamaspect/StreamSliceSupplier.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/main/scala/kafka/log/streamaspect/cache/FileCache.java b/core/src/main/scala/kafka/log/streamaspect/cache/FileCache.java index 5b02f31333..643ddfacf4 100644 --- a/core/src/main/scala/kafka/log/streamaspect/cache/FileCache.java +++ b/core/src/main/scala/kafka/log/streamaspect/cache/FileCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect.cache; diff --git a/core/src/main/scala/kafka/log/streamaspect/client/ClientFactoryProxy.java b/core/src/main/scala/kafka/log/streamaspect/client/ClientFactoryProxy.java index 6a28bb7b12..1f4d333663 100644 --- a/core/src/main/scala/kafka/log/streamaspect/client/ClientFactoryProxy.java +++ b/core/src/main/scala/kafka/log/streamaspect/client/ClientFactoryProxy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect.client; @@ -17,12 +25,12 @@ public class ClientFactoryProxy { private static final String PROTOCOL_SEPARATOR = ":"; - private static final String FACTORY_CLASS_FORMAT = "kafka.log.streamaspect.client.%s.ClientFactory"; public static Client get(Context context) { String endpoint = context.config.elasticStreamEndpoint(); String protocol = endpoint.split(PROTOCOL_SEPARATOR)[0]; - String className = String.format(FACTORY_CLASS_FORMAT, protocol); + String proxyPackage = ClientFactoryProxy.class.getPackage().getName(); + String className = String.format("%s.%s.ClientFactory", proxyPackage, protocol); try { Class clazz = Class.forName(className); Method method = clazz.getDeclaredMethod("get", Context.class); diff --git a/core/src/main/scala/kafka/log/streamaspect/client/Context.java b/core/src/main/scala/kafka/log/streamaspect/client/Context.java index 1f49743f7c..9d3b914807 100644 --- a/core/src/main/scala/kafka/log/streamaspect/client/Context.java +++ b/core/src/main/scala/kafka/log/streamaspect/client/Context.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect.client; diff --git a/core/src/main/scala/kafka/log/streamaspect/client/memory/ClientFactory.java b/core/src/main/scala/kafka/log/streamaspect/client/memory/ClientFactory.java index a028b3f3cd..4c049d71cf 100644 --- a/core/src/main/scala/kafka/log/streamaspect/client/memory/ClientFactory.java +++ b/core/src/main/scala/kafka/log/streamaspect/client/memory/ClientFactory.java @@ -1,17 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect.client.memory; -import kafka.log.streamaspect.AlwaysSuccessClient; +import kafka.log.streamaspect.ClientWrapper; import kafka.log.streamaspect.MemoryClient; import kafka.log.streamaspect.client.Context; @@ -20,7 +28,7 @@ public class ClientFactory { public static Client get(Context context) { - return new AlwaysSuccessClient(new MemoryClient()); + return new ClientWrapper(new MemoryClient()); } } diff --git a/core/src/main/scala/kafka/log/streamaspect/client/s3/ClientFactory.java b/core/src/main/scala/kafka/log/streamaspect/client/s3/ClientFactory.java index d31ef064f4..d3d52104eb 100644 --- a/core/src/main/scala/kafka/log/streamaspect/client/s3/ClientFactory.java +++ b/core/src/main/scala/kafka/log/streamaspect/client/s3/ClientFactory.java @@ -1,19 +1,27 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect.client.s3; import kafka.log.stream.s3.ConfigUtils; import kafka.log.stream.s3.DefaultS3Client; -import kafka.log.streamaspect.AlwaysSuccessClient; +import kafka.log.streamaspect.ClientWrapper; import kafka.log.streamaspect.client.ClientFactoryProxy; import kafka.log.streamaspect.client.Context; @@ -27,10 +35,13 @@ public class ClientFactory { */ public static Client get(Context context) { Config config = ConfigUtils.to(context.config); - config.nodeEpoch(System.currentTimeMillis()); config.version(() -> context.brokerServer.metadataCache().autoMQVersion().s3streamVersion()); + boolean zeroZoneChannelsEnabled = context.config.automq().zoneRouterChannels().isPresent(); + if (zeroZoneChannelsEnabled) { + config.snapshotReadEnable(true); + } DefaultS3Client client = new DefaultS3Client(context.brokerServer, config); - return new AlwaysSuccessClient(client); + return new ClientWrapper(client); } } diff --git a/core/src/main/scala/kafka/log/streamaspect/utils/ExceptionUtil.java b/core/src/main/scala/kafka/log/streamaspect/utils/ExceptionUtil.java index f14dba9b4d..fe06e66f6f 100644 --- a/core/src/main/scala/kafka/log/streamaspect/utils/ExceptionUtil.java +++ b/core/src/main/scala/kafka/log/streamaspect/utils/ExceptionUtil.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect.utils; diff --git a/core/src/main/scala/kafka/network/RequestChannel.scala b/core/src/main/scala/kafka/network/RequestChannel.scala index 3ec88d807a..365dd2e004 100644 --- a/core/src/main/scala/kafka/network/RequestChannel.scala +++ b/core/src/main/scala/kafka/network/RequestChannel.scala @@ -50,6 +50,10 @@ object RequestChannel extends Logging { private val ResponseQueueSizeMetric = "ResponseQueueSize" val ProcessorMetricTag = "processor" + // AutoMQ inject start + private val AvailableRequestSizeMetric = "AvailableRequestSize" + // AutoMQ inject end + private def isRequestLoggingEnabled: Boolean = requestLogger.underlying.isDebugEnabled sealed trait BaseRequest @@ -347,6 +351,9 @@ object RequestChannel extends Logging { } class RequestChannel(val queueSize: Int, + // AutoMQ inject start + val queuedRequestSize: Int, + // AutoMQ inject end val metricNamePrefix: String, time: Time, val metrics: RequestChannel.Metrics) { @@ -355,12 +362,23 @@ class RequestChannel(val queueSize: Int, private val metricsGroup = new KafkaMetricsGroup(this.getClass) private val requestQueue = new ArrayBlockingQueue[BaseRequest](queueSize) + // AutoMQ inject start + /** + * Queue of requests to be handled, in the order they arrived. + * Note: Before any request enters this queue, it needs to acquire {@link queuedRequestSizeSemaphore} + */ private val multiRequestQueue = new java.util.ArrayList[ArrayBlockingQueue[BaseRequest]]() - + /** + * Semaphore to limit the total size of requests in the {@link multiRequestQueue}. + */ + private val queuedRequestSizeSemaphore = new Semaphore(queuedRequestSize) + private val availableRequestSizeMetricName = metricNamePrefix.concat(AvailableRequestSizeMetric) + // AutoMQ inject end private val processors = new ConcurrentHashMap[Int, Processor]() private val requestQueueSizeMetricName = metricNamePrefix.concat(RequestQueueSizeMetric) private val responseQueueSizeMetricName = metricNamePrefix.concat(ResponseQueueSizeMetric) private val callbackQueue = new ArrayBlockingQueue[BaseRequest](queueSize) + // AutoMQ inject start private val multiCallbackQueue = new java.util.ArrayList[ArrayBlockingQueue[BaseRequest]]() private var notifiedShutdown = false @@ -371,6 +389,14 @@ class RequestChannel(val queueSize: Int, requestQueue.size() } }) + metricsGroup.newGauge(availableRequestSizeMetricName, () => { + queuedRequestSizeSemaphore.availablePermits() + }) + + def this(queueSize: Int, metricNamePrefix: String, time: Time, metrics: RequestChannel.Metrics) { + this(queueSize, Integer.MAX_VALUE, metricNamePrefix, time, metrics) + } + // AutoMQ inject end metricsGroup.newGauge(responseQueueSizeMetricName, () => { processors.values.asScala.foldLeft(0) {(total, processor) => @@ -386,14 +412,16 @@ class RequestChannel(val queueSize: Int, Map(ProcessorMetricTag -> processor.id.toString).asJava) } - def registerNRequestHandler(count: Int): util.List[BlockingQueue[BaseRequest]] = { + // AutoMQ inject start + def registerNRequestHandler(count: Int): Unit = { val queueSize = math.max(this.queueSize / count, 1) - for (i <- 0 until count) { + for (_ <- 0 until count) { multiRequestQueue.add(new ArrayBlockingQueue[BaseRequest](queueSize)) multiCallbackQueue.add(new ArrayBlockingQueue[BaseRequest](queueSize)) } Collections.unmodifiableList(multiRequestQueue) } + // AutoMQ inject end def removeProcessor(processorId: Int): Unit = { processors.remove(processorId) @@ -403,6 +431,7 @@ class RequestChannel(val queueSize: Int, /** Send a request to be handled, potentially blocking until there is room in the queue for the request */ def sendRequest(request: RequestChannel.Request): Unit = { if (multiRequestQueue.size() != 0) { + queuedRequestSizeSemaphore.acquire(Math.min(request.sizeInBytes, queuedRequestSize)) val requestQueue = multiRequestQueue.get(math.abs(request.context.connectionId.hashCode % multiRequestQueue.size())) requestQueue.put(request) } else { @@ -505,6 +534,7 @@ class RequestChannel(val queueSize: Int, } } + // AutoMQ inject start def receiveRequest(timeout: Long, id: Int): RequestChannel.BaseRequest = { val callbackQueue = multiCallbackQueue.get(id) val requestQueue = multiRequestQueue.get(id) @@ -515,11 +545,14 @@ class RequestChannel(val queueSize: Int, val request = requestQueue.poll(timeout, TimeUnit.MILLISECONDS) request match { case WakeupRequest => callbackQueue.poll() + case request: Request => + queuedRequestSizeSemaphore.release(Math.min(request.sizeInBytes, queuedRequestSize)) + request case _ => request } } } - + // AutoMQ inject end /** Get the next request or block until there is one */ @Deprecated def receiveRequest(): RequestChannel.BaseRequest = diff --git a/core/src/main/scala/kafka/network/RequestConvertToJson.scala b/core/src/main/scala/kafka/network/RequestConvertToJson.scala index 1d4366a77c..8ba7ffdb00 100644 --- a/core/src/main/scala/kafka/network/RequestConvertToJson.scala +++ b/core/src/main/scala/kafka/network/RequestConvertToJson.scala @@ -130,6 +130,7 @@ object RequestConvertToJson { case req: AutomqRegisterNodeRequest => AutomqRegisterNodeRequestDataJsonConverter.write(req.data, request.version) case req: AutomqGetNodesRequest => AutomqGetNodesRequestDataJsonConverter.write(req.data, request.version) case req: AutomqZoneRouterRequest => AutomqZoneRouterRequestDataJsonConverter.write(req.data, request.version) + case req: AutomqGetPartitionSnapshotRequest => AutomqGetPartitionSnapshotRequestDataJsonConverter.write(req.data, request.version) case req: GetNextNodeIdRequest => GetNextNodeIdRequestDataJsonConverter.write(req.data, request.version) case req: DescribeStreamsRequest => DescribeStreamsRequestDataJsonConverter.write(req.data, request.version) // AutoMQ for Kafka inject end diff --git a/core/src/main/scala/kafka/network/SocketServer.scala b/core/src/main/scala/kafka/network/SocketServer.scala index cc9018e171..67fd7559c3 100644 --- a/core/src/main/scala/kafka/network/SocketServer.scala +++ b/core/src/main/scala/kafka/network/SocketServer.scala @@ -97,7 +97,10 @@ class SocketServer(val config: KafkaConfig, private val memoryPool = if (config.queuedMaxBytes > 0) new SimpleMemoryPool(config.queuedMaxBytes, config.socketRequestMaxBytes, false, memoryPoolSensor) else MemoryPool.NONE // data-plane private[network] val dataPlaneAcceptors = new ConcurrentHashMap[EndPoint, DataPlaneAcceptor]() - val dataPlaneRequestChannel = new RequestChannel(maxQueuedRequests, DataPlaneAcceptor.MetricPrefix, time, apiVersionManager.newRequestMetrics) + // AutoMQ inject start + // val dataPlaneRequestChannel = new RequestChannel(maxQueuedRequests, DataPlaneAcceptor.MetricPrefix, time, apiVersionManager.newRequestMetrics) + val dataPlaneRequestChannel = new RequestChannel(maxQueuedRequests, config.queuedMaxRequestSize, DataPlaneAcceptor.MetricPrefix, time, apiVersionManager.newRequestMetrics) + // AutoMQ inject end // control-plane private[network] var controlPlaneAcceptorOpt: Option[ControlPlaneAcceptor] = None val controlPlaneRequestChannelOpt: Option[RequestChannel] = config.controlPlaneListenerName.map(_ => @@ -411,6 +414,10 @@ object SocketServer { val ListenerReconfigurableConfigs: Set[String] = Set(SocketServerConfigs.MAX_CONNECTIONS_CONFIG, SocketServerConfigs.MAX_CONNECTION_CREATION_RATE_CONFIG) + // AutoMQ inject start + val MaxInflightRequestsPerConnection = 64; + // AutoMQ inject end + def closeSocket( channel: SocketChannel, logging: Logging @@ -707,7 +714,10 @@ private[kafka] abstract class Acceptor(val socketServer: SocketServer, private def closeAll(): Unit = { debug("Closing server socket, selector, and any throttled sockets.") - CoreUtils.swallow(serverChannel.close(), this, Level.ERROR) + // AutoMQ inject start + // Ensure null-safe closure of serverChannel to avoid NullPointerException + Option(serverChannel).foreach { ch => CoreUtils.swallow(ch.close(), this, Level.ERROR) } + // AutoMQ inject end CoreUtils.swallow(nioSelector.close(), this, Level.ERROR) throttledSockets.foreach(throttledSocket => closeSocket(throttledSocket.socket, this)) throttledSockets.clear() @@ -1061,11 +1071,25 @@ private[kafka] class Processor( trace(s"Socket server received empty response to send, registering for read: $response") } // AutoMQ for Kafka inject start - // Try unmuting the channel. If there was no quota violation and the channel has not been throttled, - // it will be unmuted immediately. If the channel has been throttled, it will be unmuted only if the - // throttling delay has already passed by now. -// handleChannelMuteEvent(channelId, ChannelMuteEvent.RESPONSE_SENT) -// tryUnmuteChannel(channelId) + val channelContext = channelContexts.get(channelId) + val channel = selector.channel(channelId) + try{ + if (channel != null && channel.isMuted) { + val unmute = if (channelContext == null) { + true + } else if (channelContext.nextCorrelationId.size() < MaxInflightRequestsPerConnection && channelContext.clearQueueFull()) { + true + } else { + false + } + if (unmute) { + selector.unmute(channel.id) + } + } + }catch { + case e:IllegalStateException => + warn(s"Channel already closed while processing response for $channelId",e) + } // AutoMQ for Kafka inject end case response: SendResponse => @@ -1180,7 +1204,7 @@ private[kafka] class Processor( // AutoMQ will pipeline the requests to accelerate the performance and also keep the request order. // Mute the channel if the inflight requests exceed the threshold. - if (channelContext.nextCorrelationId.size() >= 8 && !channel.isMuted) { + if (channelContext.nextCorrelationId.size() >= MaxInflightRequestsPerConnection && !channel.isMuted) { if (isTraceEnabled) { trace(s"Mute channel ${channel.id} because the inflight requests exceed the threshold, inflight count is ${channelContext.nextCorrelationId.size()}.") } @@ -1228,7 +1252,7 @@ private[kafka] class Processor( if (channel.isMuted) { val unmute = if (channelContext == null) { true - } else if (channelContext.nextCorrelationId.size() < 8 && channelContext.clearQueueFull()) { + } else if (channelContext.nextCorrelationId.size() < MaxInflightRequestsPerConnection && channelContext.clearQueueFull()) { if (isTraceEnabled) { trace(s"Unmute channel ${send.destinationId} because the inflight requests are below the threshold.") } @@ -1269,7 +1293,7 @@ private[kafka] class Processor( } remove }) - channelContexts.remove(connectionId) + removeChannelContext(connectionId) // the channel has been closed by the selector but the quotas still need to be updated connectionQuotas.dec(listenerName, InetAddress.getByName(remoteHost)) } catch { @@ -1310,7 +1334,7 @@ private[kafka] class Processor( } remove }) - channelContexts.remove(connectionId) + removeChannelContext(connectionId) // inflightResponses.remove(connectionId).foreach(updateRequestMetrics) // AutoMQ for Kafka inject end } @@ -1422,6 +1446,26 @@ private[kafka] class Processor( // AutoMQ for Kafka inject end } + // AutoMQ inject start + private def removeChannelContext(connectionId: String): Unit = { + val channelContext = channelContexts.remove(connectionId) + if (channelContext == null) { + return + } + channelContext.synchronized { + channelContext.nextCorrelationId.clear() + channelContext.responses.forEach((_, response) => { + response match { + case sendResponse: SendResponse => + sendResponse.responseSend.release() + case _ => + } + }) + channelContext.responses.clear() + } + } + // AutoMQ inject end + private def dequeueResponse(): RequestChannel.Response = { val response = responseQueue.poll() if (response != null) diff --git a/core/src/main/scala/kafka/server/AlterPartitionManager.scala b/core/src/main/scala/kafka/server/AlterPartitionManager.scala index 80eba5e032..de25220979 100644 --- a/core/src/main/scala/kafka/server/AlterPartitionManager.scala +++ b/core/src/main/scala/kafka/server/AlterPartitionManager.scala @@ -403,7 +403,9 @@ class DefaultAlterPartitionManager( return } val request = new ElectLeadersRequest.Builder(ElectionType.PREFERRED, topicPartitions, 1000) - debug(s"sending elect leader to controller $request") + if (isDebugEnabled) { + debug(s"sending elect leader to controller $request") + } controllerChannelManager.sendRequest(request, new ControllerRequestCompletionHandler { override def onTimeout(): Unit = { inflightElectLeadersRequest.set(false) @@ -413,7 +415,9 @@ class DefaultAlterPartitionManager( } override def onComplete(response: ClientResponse): Unit = { - debug(s"Received elect leader response $response") + if (isDebugEnabled) { + debug(s"Received elect leader response $response") + } // no need retry, controller have backup logic to elect leader when timeout // In the normal case, check for pending updates to send immediately inflightElectLeadersRequest.set(false) diff --git a/core/src/main/scala/kafka/server/AutoTopicCreationManager.scala b/core/src/main/scala/kafka/server/AutoTopicCreationManager.scala index a5c0146f84..9060c464bf 100644 --- a/core/src/main/scala/kafka/server/AutoTopicCreationManager.scala +++ b/core/src/main/scala/kafka/server/AutoTopicCreationManager.scala @@ -24,9 +24,10 @@ import kafka.controller.KafkaController import kafka.coordinator.transaction.TransactionCoordinator import kafka.utils.Logging import org.apache.kafka.clients.ClientResponse +import org.apache.kafka.common.config.TopicConfig import org.apache.kafka.common.errors.InvalidTopicException import org.apache.kafka.common.internals.Topic -import org.apache.kafka.common.internals.Topic.{GROUP_METADATA_TOPIC_NAME, TRANSACTION_STATE_TOPIC_NAME} +import org.apache.kafka.common.internals.Topic.{GROUP_METADATA_TOPIC_NAME, TABLE_TOPIC_CONTROL_TOPIC_NAME, TABLE_TOPIC_DATA_TOPIC_NAME, TRANSACTION_STATE_TOPIC_NAME} import org.apache.kafka.common.message.CreateTopicsRequestData import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicConfig, CreatableTopicConfigCollection} import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic @@ -244,6 +245,28 @@ class DefaultAutoTopicCreationManager( .setReplicationFactor(config.transactionTopicReplicationFactor) .setConfigs(convertToTopicConfigCollections( txnCoordinator.transactionTopicConfigs)) + + // AutoMQ inject start + case TABLE_TOPIC_CONTROL_TOPIC_NAME => { + val configs = new Properties() + configs.put(TopicConfig.MAX_MESSAGE_BYTES_CONFIG, 20 * 1024 * 1024) + new CreatableTopic() + .setName(topic) + .setNumPartitions(1) + .setReplicationFactor(1) + .setConfigs(convertToTopicConfigCollections(configs)) + } + case TABLE_TOPIC_DATA_TOPIC_NAME => { + val configs = new Properties() + configs.put(TopicConfig.MAX_MESSAGE_BYTES_CONFIG, 20 * 1024 * 1024) + new CreatableTopic() + .setName(topic) + .setNumPartitions(50) + .setReplicationFactor(1) + .setConfigs(convertToTopicConfigCollections(configs)) + } + // AutoMQ inject end + case topicName => new CreatableTopic() .setName(topicName) diff --git a/core/src/main/scala/kafka/server/BrokerServer.scala b/core/src/main/scala/kafka/server/BrokerServer.scala index feccbfff7c..534854fa48 100644 --- a/core/src/main/scala/kafka/server/BrokerServer.scala +++ b/core/src/main/scala/kafka/server/BrokerServer.scala @@ -17,7 +17,13 @@ package kafka.server -import kafka.automq.zonerouter.{NoopProduceRouter, ProduceRouter} +import com.automq.stream.s3.S3Storage +import kafka.automq.backpressure.{BackPressureConfig, BackPressureManager, DefaultBackPressureManager, Regulator} +import kafka.automq.failover.FailoverListener +import kafka.automq.kafkalinking.KafkaLinkingManager +import kafka.automq.interceptor.{NoopTrafficInterceptor, TrafficInterceptor} +import kafka.automq.table.TableManager +import kafka.automq.zerozone.{ConfirmWALProvider, DefaultClientRackProvider, DefaultConfirmWALProvider, DefaultRouterChannelProvider, DefaultLinkRecordDecoder, RouterChannelProvider, ZeroZoneTrafficInterceptor} import kafka.cluster.EndPoint import kafka.coordinator.group.{CoordinatorLoaderImpl, CoordinatorPartitionWriter, GroupCoordinatorAdapter} import kafka.coordinator.transaction.{ProducerIdManager, TransactionCoordinator} @@ -27,7 +33,7 @@ import kafka.log.streamaspect.ElasticLogManager import kafka.network.{DataPlaneAcceptor, SocketServer} import kafka.raft.KafkaRaftManager import kafka.server.metadata.{AclPublisher, BrokerMetadataPublisher, ClientQuotaMetadataManager, DelegationTokenPublisher, DynamicClientQuotaPublisher, DynamicConfigPublisher, KRaftMetadataCache, ScramPublisher} -import kafka.server.streamaspect.{ElasticKafkaApis, ElasticReplicaManager} +import kafka.server.streamaspect.{ElasticKafkaApis, ElasticReplicaManager, PartitionLifecycleListener} import kafka.utils.CoreUtils import org.apache.kafka.common.config.ConfigException import org.apache.kafka.common.feature.SupportedVersionRange @@ -39,7 +45,7 @@ import org.apache.kafka.common.security.token.delegation.internals.DelegationTok import org.apache.kafka.common.utils.{LogContext, Time} import org.apache.kafka.common.{ClusterResource, TopicPartition, Uuid} import org.apache.kafka.coordinator.group.metrics.{GroupCoordinatorMetrics, GroupCoordinatorRuntimeMetrics} -import org.apache.kafka.coordinator.group.{CoordinatorRecord, GroupCoordinator, GroupCoordinatorService, CoordinatorRecordSerde} +import org.apache.kafka.coordinator.group.{CoordinatorRecord, CoordinatorRecordSerde, GroupCoordinator, GroupCoordinatorService} import org.apache.kafka.image.publisher.{BrokerRegistrationTracker, MetadataPublisher} import org.apache.kafka.image.loader.MetadataLoader import org.apache.kafka.metadata.{BrokerState, ListenerInfo, VersionRange} @@ -158,7 +164,17 @@ class BrokerServer( def metadataLoader: MetadataLoader = sharedServer.loader - var produceRouter: ProduceRouter = _ + var routerChannelProvider: RouterChannelProvider = _ + var confirmWALProvider: ConfirmWALProvider = _ + var trafficInterceptor: TrafficInterceptor = _ + + var backPressureManager: BackPressureManager = _ + + val clientRackProvider = new DefaultClientRackProvider() + // init reconfigurable before startup + config.addReconfigurable(clientRackProvider) + + var tableManager: TableManager = _ // AutoMQ inject end private def maybeChangeStatus(from: ProcessStatus, to: ProcessStatus): Boolean = { @@ -355,6 +371,11 @@ class BrokerServer( groupCoordinator = createGroupCoordinator() + // AutoMQ injection start + this._replicaManager.setKafkaLinkingManager(newKafkaLinkingManager()) + groupCoordinator = createGroupCoordinatorWrapper(groupCoordinator) + // AutoMQ injection end + val producerIdManagerSupplier = () => ProducerIdManager.rpc( config.brokerId, time, @@ -508,6 +529,13 @@ class BrokerServer( () => lifecycleManager.resendBrokerRegistrationUnlessZkMode()) metadataPublishers.add(brokerRegistrationTracker) + // AutoMQ inject start + backPressureManager = new DefaultBackPressureManager( + BackPressureConfig.from(config), + newBackPressureRegulator() + ) + backPressureManager.start() + // AutoMQ inject end // Register parts of the broker that can be reconfigured via dynamic configs. This needs to // be done before we publish the dynamic configs, so that we don't miss anything. @@ -542,8 +570,23 @@ class BrokerServer( // AutoMQ inject start + routerChannelProvider = newRouterChannelProvider() + confirmWALProvider = newConfirmWALProvider() + if (routerChannelProvider != null) { + S3Storage.setLinkRecordDecoder(new DefaultLinkRecordDecoder(routerChannelProvider)) + } ElasticLogManager.init(config, clusterId, this) - produceRouter = newProduceRouter() + trafficInterceptor = newTrafficInterceptor() + dataPlaneRequestProcessor.asInstanceOf[ElasticKafkaApis].setTrafficInterceptor(trafficInterceptor) + replicaManager.setTrafficInterceptor(trafficInterceptor) + replicaManager.setS3StreamContext(com.automq.stream.Context.instance()) + + tableManager = new TableManager(metadataCache, config) + newPartitionLifecycleListeners().forEach(l => { + _replicaManager.addPartitionLifecycleListener(l) + }) + + newFailoverListener(ElasticLogManager.INSTANCE.get.client) // AutoMQ inject end // We're now ready to unfence the broker. This also allows this broker to transition @@ -625,6 +668,10 @@ class BrokerServer( } } + protected def createGroupCoordinatorWrapper(groupCoordinator: GroupCoordinator): GroupCoordinator = { + groupCoordinator + } + protected def createRemoteLogManager(): Option[RemoteLogManager] = { if (config.remoteLogManagerConfig.isRemoteStorageSystemEnabled()) { Some(new RemoteLogManager(config.remoteLogManagerConfig, config.brokerId, config.logDirs.head, clusterId, time, @@ -664,12 +711,21 @@ class BrokerServer( lifecycleManager.beginShutdown() // AutoMQ for Kafka inject start + if (backPressureManager != null) { + CoreUtils.swallow(backPressureManager.shutdown(), this) + } + // https://github.com/AutoMQ/automq-for-kafka/issues/540 // await partition shutdown: // 1. after lifecycleManager start shutdown to trigger partitions gracefully reassign. // 2. before metadataListener start close to ensure S3Stream can read the latest metadata. if (replicaManager != null) { CoreUtils.swallow(replicaManager.awaitAllPartitionShutdown(), this) + CoreUtils.swallow(trafficInterceptor.close(), this) + if (routerChannelProvider != null) { + CoreUtils.swallow(routerChannelProvider.close(), this) + } + CoreUtils.swallow(ElasticLogManager.shutdown(), this) } // AutoMQ for Kafka inject end @@ -789,10 +845,56 @@ class BrokerServer( ) } - protected def newProduceRouter(): ProduceRouter = { - val produceRouter = new NoopProduceRouter(dataPlaneRequestProcessor.asInstanceOf[ElasticKafkaApis], metadataCache) - dataPlaneRequestProcessor.asInstanceOf[ElasticKafkaApis].setProduceRouter(produceRouter) - produceRouter + protected def newRouterChannelProvider(): RouterChannelProvider = { + if (config.automq.zoneRouterChannels().isEmpty) { + return null + } + val bucketURI = config.automq.zoneRouterChannels.get.get(0) + new DefaultRouterChannelProvider(config.nodeId, config.automq.nodeEpoch, bucketURI, dataPlaneRequestProcessor.clusterId) + } + + protected def newConfirmWALProvider(): ConfirmWALProvider = { + if (config.automq.zoneRouterChannels().isEmpty) { + return null + } + new DefaultConfirmWALProvider(dataPlaneRequestProcessor.clusterId) + } + + protected def newTrafficInterceptor(): TrafficInterceptor = { + val trafficInterceptor = if (config.automq.zoneRouterChannels().isEmpty) { + new NoopTrafficInterceptor(dataPlaneRequestProcessor.asInstanceOf[ElasticKafkaApis], metadataCache) + } else { + val zeroZoneRouter = new ZeroZoneTrafficInterceptor(routerChannelProvider, confirmWALProvider, dataPlaneRequestProcessor.asInstanceOf[ElasticKafkaApis], metadataCache, clientRackProvider, config) + metadataLoader.installPublishers(util.List.of(zeroZoneRouter)) + zeroZoneRouter + } + trafficInterceptor + } + + protected def newPartitionLifecycleListeners(): util.List[PartitionLifecycleListener] = { + val list = new util.ArrayList[PartitionLifecycleListener]() + list.add(tableManager) + list + } + + protected def newBackPressureRegulator(): Regulator = { + new Regulator { + override def increase(): Unit = { + } + + override def decrease(): Unit = { + } + } + } + + protected def newKafkaLinkingManager(): KafkaLinkingManager = { + null + } + + protected def newFailoverListener(client: com.automq.stream.api.Client): FailoverListener = { + val failoverListener = new FailoverListener(config.nodeId, client) + metadataLoader.installPublishers(util.List.of(failoverListener)); + failoverListener } // AutoMQ inject end diff --git a/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala b/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala index 997e6c1859..8974248131 100644 --- a/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala +++ b/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala @@ -44,7 +44,7 @@ class ClientRequestQuotaManager(private val config: ClientQuotaManagerConfig, private val quotaCallback: Option[ClientQuotaCallback]) extends ClientQuotaManager(config, metrics, QuotaType.Request, time, threadNamePrefix, quotaCallback) { - protected val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(this.config.quotaWindowSizeSeconds) + private val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(this.config.quotaWindowSizeSeconds) private val exemptMetricName = metrics.metricName("exempt-request-time", QuotaType.Request.toString, "Tracking exempt-request-time utilization percentage") diff --git a/core/src/main/scala/kafka/server/ConfigHandler.scala b/core/src/main/scala/kafka/server/ConfigHandler.scala index d21e6d9dc0..44b49a1189 100644 --- a/core/src/main/scala/kafka/server/ConfigHandler.scala +++ b/core/src/main/scala/kafka/server/ConfigHandler.scala @@ -237,15 +237,15 @@ class IpConfigHandler(private val connectionQuotas: ConnectionQuotas) extends Co class BrokerConfigHandler(private val brokerConfig: KafkaConfig, private val quotaManagers: QuotaManagers) extends ConfigHandler with Logging { def processConfigChanges(brokerId: String, properties: Properties): Unit = { + // AutoMQ for Kafka inject start + if (brokerId == ZooKeeperInternals.DEFAULT_STRING || brokerConfig.brokerId == brokerId.trim.toInt) { + quotaManagers.broker.updateQuotaConfigs(Some(properties)) + } + // AutoMQ for Kafka inject end if (brokerId == ZooKeeperInternals.DEFAULT_STRING) brokerConfig.dynamicConfig.updateDefaultConfig(properties) else if (brokerConfig.brokerId == brokerId.trim.toInt) { brokerConfig.dynamicConfig.updateBrokerConfig(brokerConfig.brokerId, properties) - - // AutoMQ for Kafka inject start - quotaManagers.broker.updateQuotaConfigs(Some(properties)) - // AutoMQ for Kafka inject end - } val updatedDynamicBrokerConfigs = brokerConfig.dynamicConfig.currentDynamicBrokerConfigs val updatedDynamicDefaultConfigs = brokerConfig.dynamicConfig.currentDynamicDefaultConfigs diff --git a/core/src/main/scala/kafka/server/ControllerConfigurationValidator.scala b/core/src/main/scala/kafka/server/ControllerConfigurationValidator.scala index b99065b573..bff152695f 100644 --- a/core/src/main/scala/kafka/server/ControllerConfigurationValidator.scala +++ b/core/src/main/scala/kafka/server/ControllerConfigurationValidator.scala @@ -17,16 +17,16 @@ package kafka.server -import java.util -import java.util.Properties import org.apache.kafka.common.config.ConfigResource import org.apache.kafka.common.config.ConfigResource.Type.{BROKER, CLIENT_METRICS, TOPIC} -import org.apache.kafka.controller.ConfigurationValidator import org.apache.kafka.common.errors.{InvalidConfigurationException, InvalidRequestException} import org.apache.kafka.common.internals.Topic +import org.apache.kafka.controller.ConfigurationValidator import org.apache.kafka.server.metrics.ClientMetricsConfigs import org.apache.kafka.storage.internals.log.LogConfig +import java.util +import java.util.Properties import scala.collection.mutable /** @@ -109,6 +109,11 @@ class ControllerConfigurationValidator(kafkaConfig: KafkaConfig) extends Configu } LogConfig.validate(properties, kafkaConfig.extractLogConfigMap, kafkaConfig.remoteLogManagerConfig.isRemoteStorageSystemEnabled()) + + // AutoMQ inject start + LogConfig.validateTableTopicSchemaConfigValues(properties, kafkaConfig.tableTopicSchemaRegistryUrl) + // AutoMQ inject end + case BROKER => validateBrokerName(resource.name()) case CLIENT_METRICS => val properties = new Properties() diff --git a/core/src/main/scala/kafka/server/ControllerServer.scala b/core/src/main/scala/kafka/server/ControllerServer.scala index af2c9a98fe..099acd1d13 100644 --- a/core/src/main/scala/kafka/server/ControllerServer.scala +++ b/core/src/main/scala/kafka/server/ControllerServer.scala @@ -21,15 +21,13 @@ import com.automq.stream.s3.Constants import com.automq.stream.s3.metadata.ObjectUtils import kafka.autobalancer.AutoBalancerManager import kafka.autobalancer.services.AutoBalancerService +import kafka.automq.controller.DefaultQuorumControllerExtension import kafka.controller.streamaspect.client.{Context, StreamClientFactoryProxy} - import kafka.migration.MigrationPropagator import kafka.network.{DataPlaneAcceptor, SocketServer} import kafka.raft.KafkaRaftManager import kafka.server.QuotaFactory.QuotaManagers - -import scala.collection.immutable -import kafka.server.metadata.{AclPublisher, ClientQuotaMetadataManager, DelegationTokenPublisher, DynamicClientQuotaPublisher, DynamicConfigPublisher, KRaftMetadataCache, KRaftMetadataCachePublisher, ScramPublisher} +import kafka.server.metadata._ import kafka.server.streamaspect.ElasticControllerApis import kafka.utils.{CoreUtils, Logging} import kafka.zk.{KafkaZkClient, ZkMigrationClient} @@ -38,32 +36,33 @@ import org.apache.kafka.common.network.ListenerName import org.apache.kafka.common.security.scram.internals.ScramMechanism import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache import org.apache.kafka.common.utils.LogContext -import org.apache.kafka.common.{ClusterResource, Endpoint, Uuid} +import org.apache.kafka.common.{ClusterResource, Endpoint, Reconfigurable, Uuid} import org.apache.kafka.controller.metrics.{ControllerMetadataMetricsPublisher, QuorumControllerMetrics} import org.apache.kafka.controller.{QuorumController, QuorumControllerExtension, QuorumFeatures} import org.apache.kafka.image.publisher.{ControllerRegistrationsPublisher, MetadataPublisher} -import org.apache.kafka.metadata.{KafkaConfigSchema, ListenerInfo} import org.apache.kafka.metadata.authorizer.ClusterMetadataAuthorizer import org.apache.kafka.metadata.bootstrap.BootstrapMetadata import org.apache.kafka.metadata.migration.{KRaftMigrationDriver, LegacyPropagator} import org.apache.kafka.metadata.placement.{ReplicaPlacer, StripedReplicaPlacer} import org.apache.kafka.metadata.publisher.FeaturesPublisher +import org.apache.kafka.metadata.{KafkaConfigSchema, ListenerInfo} import org.apache.kafka.raft.QuorumConfig import org.apache.kafka.security.{CredentialProvider, PasswordEncoder} import org.apache.kafka.server.NodeToControllerChannelManager import org.apache.kafka.server.authorizer.Authorizer -import org.apache.kafka.server.config.ServerLogConfigs.{ALTER_CONFIG_POLICY_CLASS_NAME_CONFIG, CREATE_TOPIC_POLICY_CLASS_NAME_CONFIG} import org.apache.kafka.server.common.ApiMessageAndVersion import org.apache.kafka.server.config.ConfigType +import org.apache.kafka.server.config.ServerLogConfigs.{ALTER_CONFIG_POLICY_CLASS_NAME_CONFIG, CREATE_TOPIC_POLICY_CLASS_NAME_CONFIG} import org.apache.kafka.server.metrics.{KafkaMetricsGroup, KafkaYammerMetrics, LinuxIoMetricsCollector} import org.apache.kafka.server.network.{EndpointReadyFutures, KafkaAuthorizerServerInfo} import org.apache.kafka.server.policy.{AlterConfigPolicy, CreateTopicPolicy} import org.apache.kafka.server.util.{Deadline, FutureUtils} import java.util -import java.util.{Optional, OptionalLong, Random} import java.util.concurrent.locks.ReentrantLock import java.util.concurrent.{CompletableFuture, TimeUnit} +import java.util.{Optional, OptionalLong, Random} +import scala.collection.immutable import scala.compat.java8.OptionConverters._ import scala.jdk.CollectionConverters._ @@ -535,8 +534,9 @@ class ControllerServer( if (socketServer != null) CoreUtils.swallow(socketServer.stopProcessingRequests(), this) migrationSupport.foreach(_.shutdown(this)) - if (controller != null) + if (controller != null) { controller.beginShutdown() + } if (socketServer != null) CoreUtils.swallow(socketServer.shutdown(), this) if (controllerApisHandlerPool != null) @@ -578,11 +578,18 @@ class ControllerServer( // AutoMQ for Kafka inject start protected def quorumControllerExtension(quorumController: QuorumController): QuorumControllerExtension = { - QuorumControllerExtension.NOOP + new DefaultQuorumControllerExtension(quorumController) } protected def replicaPlacer(): ReplicaPlacer = { new StripedReplicaPlacer(new Random()) } + + // return a list of all reconfigurable objects + def reconfigurables(): java.util.List[Reconfigurable] = { + java.util.List.of( + autoBalancerManager + ) + } // AutoMQ for Kafka inject end } diff --git a/core/src/main/scala/kafka/server/DelayedFetch.scala b/core/src/main/scala/kafka/server/DelayedFetch.scala index a3c5e64d93..6bb9d6257c 100644 --- a/core/src/main/scala/kafka/server/DelayedFetch.scala +++ b/core/src/main/scala/kafka/server/DelayedFetch.scala @@ -241,7 +241,6 @@ class DelayedFetch( error(s"Unexpected error in delayed fetch: $params $fetchInfos ", e) } } - ReadHint.clear() // AutoMQ for Kafka inject end val fetchPartitionData = logReadResults.map { case (tp, result) => @@ -252,6 +251,11 @@ class DelayedFetch( } responseCallback(fetchPartitionData) + // AutoMQ for Kafka inject start + // clear hint after callback as it will be used in the callback + // see {@link ElasticKafkaApis#handleFetchRequest#processResponseCallback} + ReadHint.clear() + // AutoMQ for Kafka inject end } } diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala index 820086c6a4..f2b84572b2 100755 --- a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala +++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala @@ -18,6 +18,7 @@ package kafka.server import kafka.autobalancer.config.{AutoBalancerControllerConfig, AutoBalancerMetricsReporterConfig} +import kafka.automq.backpressure.BackPressureConfig import java.util import java.util.{Collections, Properties} @@ -42,7 +43,7 @@ import org.apache.kafka.coordinator.transaction.TransactionLogConfigs import org.apache.kafka.network.SocketServerConfigs import org.apache.kafka.security.PasswordEncoder import org.apache.kafka.server.ProcessRole -import org.apache.kafka.server.config.{ConfigType, ServerConfigs, ReplicationConfigs, ServerLogConfigs, ServerTopicConfigSynonyms, ZooKeeperInternals} +import org.apache.kafka.server.config.{ConfigType, ReplicationConfigs, ServerConfigs, ServerLogConfigs, ServerTopicConfigSynonyms, ZooKeeperInternals} import org.apache.kafka.server.log.remote.storage.RemoteLogManagerConfig import org.apache.kafka.server.metrics.{ClientMetricsReceiverPlugin, MetricConfigs} import org.apache.kafka.server.telemetry.ClientTelemetry @@ -101,7 +102,8 @@ object DynamicBrokerConfig { DynamicProducerStateManagerConfig ++ DynamicRemoteLogConfig.ReconfigurableConfigs ++ AutoBalancerControllerConfig.RECONFIGURABLE_CONFIGS.asScala ++ - AutoBalancerMetricsReporterConfig.RECONFIGURABLE_CONFIGS.asScala + AutoBalancerMetricsReporterConfig.RECONFIGURABLE_CONFIGS.asScala ++ + BackPressureConfig.RECONFIGURABLE_CONFIGS.asScala private val ClusterLevelListenerConfigs = Set(SocketServerConfigs.MAX_CONNECTIONS_CONFIG, SocketServerConfigs.MAX_CONNECTION_CREATION_RATE_CONFIG, SocketServerConfigs.NUM_NETWORK_THREADS_CONFIG) private val PerBrokerConfigs = (DynamicSecurityConfigs ++ DynamicListenerConfig.ReconfigurableConfigs).diff( @@ -269,6 +271,13 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging addReconfigurable(kafkaServer.kafkaYammerMetrics) addReconfigurable(new DynamicMetricsReporters(kafkaConfig.brokerId, kafkaServer.config, kafkaServer.metrics, kafkaServer.clusterId)) addReconfigurable(new DynamicClientQuotaCallback(kafkaServer.quotaManagers, kafkaServer.config)) + // AutoMQ inject start + kafkaServer match { + case brokerServer: BrokerServer => + addReconfigurable(brokerServer.backPressureManager) + case _ => + } + // AutoMQ inject end addBrokerReconfigurable(new BrokerDynamicThreadPool(kafkaServer)) addBrokerReconfigurable(new DynamicLogConfig(kafkaServer.logManager, kafkaServer)) @@ -293,7 +302,9 @@ class DynamicBrokerConfig(private val kafkaConfig: KafkaConfig) extends Logging addReconfigurable(new DynamicMetricsReporters(kafkaConfig.nodeId, controller.config, controller.metrics, controller.clusterId)) } addReconfigurable(new DynamicClientQuotaCallback(controller.quotaManagers, controller.config)) - addReconfigurable(controller.autoBalancerManager) + // AutoMQ inject start + controller.reconfigurables().asScala.foreach(addReconfigurable) + // AutoMQ inject end addBrokerReconfigurable(new ControllerDynamicThreadPool(controller)) // TODO: addBrokerReconfigurable(new DynamicListenerConfig(controller)) addBrokerReconfigurable(controller.socketServer) diff --git a/core/src/main/scala/kafka/server/FairLimiter.java b/core/src/main/scala/kafka/server/FairLimiter.java index d6ad1edc6c..273c64025d 100644 --- a/core/src/main/scala/kafka/server/FairLimiter.java +++ b/core/src/main/scala/kafka/server/FairLimiter.java @@ -1,65 +1,54 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.server; import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; -import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.ReentrantLock; /** * A fair limiter whose {@link #acquire} method is fair, i.e. the waiting threads are served in the order of arrival. */ public class FairLimiter implements Limiter { private final int maxPermits; + private final Semaphore permits; + /** - * The lock used to protect @{link #acquireLocked} + * The name of this limiter, used for metrics. */ - private final Lock lock = new ReentrantLock(true); - private final Semaphore permits; + private final String name; - public FairLimiter(int size) { - maxPermits = size; - permits = new Semaphore(size); + public FairLimiter(int size, String name) { + this.maxPermits = size; + this.permits = new Semaphore(size, true); + this.name = name; } @Override public Handler acquire(int permit) throws InterruptedException { - lock.lock(); - try { - permits.acquire(permit); - return new FairHandler(permit); - } finally { - lock.unlock(); - } + permits.acquire(permit); + return new FairHandler(permit); } @Override public Handler acquire(int permit, long timeoutMs) throws InterruptedException { - long start = System.nanoTime(); - if (lock.tryLock(timeoutMs, TimeUnit.MILLISECONDS)) { - try { - // calculate the time left for {@code acquireLocked} - long elapsed = System.nanoTime() - start; - long left = TimeUnit.MILLISECONDS.toNanos(timeoutMs) - elapsed; - // note: {@code left} may be negative here, but it's OK for acquireLocked - return acquireLocked(permit, left); - } finally { - lock.unlock(); - } - } else { - // tryLock timeout, return null - return null; - } + return acquireLocked(permit, timeoutMs); } @Override @@ -72,6 +61,16 @@ public int availablePermits() { return permits.availablePermits(); } + @Override + public int waitingThreads() { + return permits.getQueueLength(); + } + + @Override + public String name() { + return name; + } + private Handler acquireLocked(int permit, long timeoutNs) throws InterruptedException { if (permit > maxPermits) { permit = maxPermits; diff --git a/core/src/main/scala/kafka/server/FetchSession.scala b/core/src/main/scala/kafka/server/FetchSession.scala index 8cdcee2b66..2b0ec26476 100644 --- a/core/src/main/scala/kafka/server/FetchSession.scala +++ b/core/src/main/scala/kafka/server/FetchSession.scala @@ -22,6 +22,7 @@ import kafka.utils.Logging import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.common.message.FetchResponseData import org.apache.kafka.common.protocol.Errors +import org.apache.kafka.common.record.PooledResource import org.apache.kafka.common.requests.FetchMetadata.{FINAL_EPOCH, INITIAL_EPOCH, INVALID_SESSION_ID} import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, FetchMetadata => JFetchMetadata} import org.apache.kafka.common.utils.{ImplicitLinkedHashCollection, Time, Utils} @@ -503,6 +504,13 @@ class IncrementalFetchContext(private val time: Time, } } else { if (updateFetchContextAndRemoveUnselected) { + // AutoMQ inject start + respData.records() match { + case r: PooledResource => + r.release() + case _ => + } + // AutoMQ inject end iter.remove() } } @@ -540,6 +548,16 @@ class IncrementalFetchContext(private val time: Time, if (session.epoch != expectedEpoch) { info(s"Incremental fetch session ${session.id} expected epoch $expectedEpoch, but " + s"got ${session.epoch}. Possible duplicate request.") + // AutoMQ inject start + // The fetch will return empty data, so we need to release the fetched records. + updates.forEach((_, response) => { + response.records() match { + case r: PooledResource => + r.release() + case _ => + } + }) + // AutoMQ inject end FetchResponse.of(Errors.INVALID_FETCH_SESSION_EPOCH, 0, session.id, new FetchSession.RESP_MAP) } else { // Iterate over the update list using PartitionIterator. This will prune updates which don't need to be sent diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala index 65b65fa3d8..520b15c2bb 100644 --- a/core/src/main/scala/kafka/server/KafkaApis.scala +++ b/core/src/main/scala/kafka/server/KafkaApis.scala @@ -18,6 +18,7 @@ package kafka.server import kafka.api.ElectLeadersRequestOps +import kafka.automq.interceptor.ClientIdMetadata import kafka.controller.ReplicaAssignment import kafka.coordinator.transaction.{InitProducerIdResult, TransactionCoordinator} import kafka.network.RequestChannel @@ -70,7 +71,7 @@ import org.apache.kafka.common.{Node, TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.coordinator.group.{Group, GroupCoordinator} import org.apache.kafka.server.ClientMetricsManager import org.apache.kafka.server.authorizer._ -import org.apache.kafka.server.common.{MetadataVersion} +import org.apache.kafka.server.common.MetadataVersion import org.apache.kafka.server.common.MetadataVersion.{IBP_0_11_0_IV0, IBP_2_3_IV0} import org.apache.kafka.server.record.BrokerCompressionType import org.apache.kafka.storage.internals.log.{AppendOrigin, FetchIsolation, FetchParams, FetchPartitionData} @@ -1068,7 +1069,15 @@ class KafkaApis(val requestChannel: RequestChannel, request.context.principal, request.context.listenerName.value)) } else { - Optional.empty() + // AutoMQ inject start + // Both leader partition and snapshot-read partition can handle the fetch request. + Optional.of(new DefaultClientMetadata( + null, + clientId, + request.context.clientAddress, + request.context.principal, + request.context.listenerName.value)) + // AutoMQ inject end } val params = new FetchParams( @@ -1098,9 +1107,9 @@ class KafkaApis(val requestChannel: RequestChannel, def handleListOffsetRequest(request: RequestChannel.Request): Unit = { val version = request.header.apiVersion - val topics = if (version == 0) + val topics = if (version == 0) { handleListOffsetRequestV0(request) - else + } else handleListOffsetRequestV1AndAbove(request) requestHelper.sendResponseMaybeThrottle(request, requestThrottleMs => new ListOffsetsResponse(new ListOffsetsResponseData() @@ -1288,6 +1297,14 @@ class KafkaApis(val requestChannel: RequestChannel, val controllerMutationQuota = quotas.controllerMutation.newPermissiveQuotaFor(request) autoTopicCreationManager.createTopics(nonExistingTopics, controllerMutationQuota, Some(request.context)) } else { + // AutoMQ inject start + for (tableTopic <- Set(Topic.TABLE_TOPIC_CONTROL_TOPIC_NAME, Topic.TABLE_TOPIC_DATA_TOPIC_NAME)) { + if (nonExistingTopics.contains(tableTopic)) { + val controllerMutationQuota = quotas.controllerMutation.newPermissiveQuotaFor(request) + autoTopicCreationManager.createTopics(Set(tableTopic), controllerMutationQuota, Some(request.context)) + } + } + // AutoMQ inject end nonExistingTopics.map { topic => val error = try { Topic.validate(topic) @@ -1428,7 +1445,7 @@ class KafkaApis(val requestChannel: RequestChannel, } } - val clientId = request.header.clientId() + val clientId = ClientIdMetadata.of(request.header.clientId(), request.context.clientAddress, request.context.connectionId) val listenerName = request.context.listenerName.value() requestHelper.sendResponseMaybeThrottle(request, requestThrottleMs => MetadataResponse.prepareResponse( @@ -4112,7 +4129,7 @@ class KafkaApis(val requestChannel: RequestChannel, } // AutoMQ inject start - protected def metadataTopicsInterceptor(clientId: String, listenerName: String, topics: util.List[MetadataResponseData.MetadataResponseTopic]): util.List[MetadataResponseData.MetadataResponseTopic] = { + protected def metadataTopicsInterceptor(clientId: ClientIdMetadata, listenerName: String, topics: util.List[MetadataResponseData.MetadataResponseTopic]): util.List[MetadataResponseData.MetadataResponseTopic] = { topics } diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala index 407aa1b9eb..e15f350cf0 100755 --- a/core/src/main/scala/kafka/server/KafkaConfig.scala +++ b/core/src/main/scala/kafka/server/KafkaConfig.scala @@ -100,7 +100,7 @@ object KafkaConfig { AutoBalancerControllerConfig.CONFIG_DEF.configKeys().values().forEach(key => configDef.define(key)) AutoBalancerMetricsReporterConfig.CONFIG_DEF.configKeys().values().forEach(key => configDef.define(key)) // AutoMQ inject end - + def configNames: Seq[String] = configDef.names.asScala.toBuffer.sorted private[server] def defaultValues: Map[String, _] = configDef.defaultValues.asScala @@ -435,6 +435,9 @@ class KafkaConfig private(doLog: Boolean, val props: util.Map[_, _]) val connectionsMaxIdleMs = getLong(SocketServerConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG) val failedAuthenticationDelayMs = getInt(SocketServerConfigs.FAILED_AUTHENTICATION_DELAY_MS_CONFIG) val queuedMaxRequests = getInt(SocketServerConfigs.QUEUED_MAX_REQUESTS_CONFIG) + // AutoMQ inject start + val queuedMaxRequestSize = getInt(SocketServerConfigs.QUEUED_MAX_REQUESTS_SIZE_BYTES_CONFIG) + // AutoMQ inject end val queuedMaxBytes = getLong(SocketServerConfigs.QUEUED_MAX_BYTES_CONFIG) def numNetworkThreads = getInt(SocketServerConfigs.NUM_NETWORK_THREADS_CONFIG) @@ -772,6 +775,7 @@ class KafkaConfig private(doLog: Boolean, val props: util.Map[_, _]) val s3ObjectPartSize = getInt(AutoMQConfig.S3_OBJECT_PART_SIZE_CONFIG) val s3StreamAllocatorPolicy = Enum.valueOf(classOf[ByteBufAllocPolicy], getString(AutoMQConfig.S3_STREAM_ALLOCATOR_POLICY_CONFIG)) val (s3WALCacheSize, s3BlockCacheSize, s3WALUploadThreshold) = adjustS3Configs(s3StreamAllocatorPolicy) + val s3WALUploadIntervalMs = getLong(AutoMQConfig.S3_WAL_UPLOAD_INTERVAL_MS_CONFIG) val s3StreamObjectCompactionTaskIntervalMinutes = getInt(AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_INTERVAL_MINUTES_CONFIG) val s3StreamObjectCompactionMaxSizeBytes = getLong(AutoMQConfig.S3_STREAM_OBJECT_COMPACTION_MAX_SIZE_BYTES_CONFIG) val s3ControllerRequestRetryMaxCount = getInt(AutoMQConfig.S3_CONTROLLER_REQUEST_RETRY_MAX_COUNT_CONFIG) @@ -789,6 +793,10 @@ class KafkaConfig private(doLog: Boolean, val props: util.Map[_, _]) val s3RefillPeriodMsProp = getInt(AutoMQConfig.S3_NETWORK_REFILL_PERIOD_MS_CONFIG) val s3MetricsLevel = getString(AutoMQConfig.S3_TELEMETRY_METRICS_LEVEL_CONFIG) val s3ExporterReportIntervalMs = getInt(AutoMQConfig.S3_TELEMETRY_EXPORTER_REPORT_INTERVAL_MS_CONFIG) + val tableTopicNamespace = getString(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG) + val s3BackPressureEnabled = getBoolean(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG) + val s3BackPressureCooldownMs = getLong(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG) + val tableTopicSchemaRegistryUrl = getString(AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG) // AutoMQ inject end /** Internal Configurations **/ @@ -1228,6 +1236,16 @@ class KafkaConfig private(doLog: Boolean, val props: util.Map[_, _]) logProps.put(TopicConfig.MESSAGE_DOWNCONVERSION_ENABLE_CONFIG, logMessageDownConversionEnable: java.lang.Boolean) logProps.put(TopicConfig.LOCAL_LOG_RETENTION_MS_CONFIG, remoteLogManagerConfig.logLocalRetentionMs: java.lang.Long) logProps.put(TopicConfig.LOCAL_LOG_RETENTION_BYTES_CONFIG, remoteLogManagerConfig.logLocalRetentionBytes: java.lang.Long) + + // AutoMQ inject start + if (tableTopicNamespace != null) { + logProps.put(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG, tableTopicNamespace) + } + if (tableTopicSchemaRegistryUrl != null) { + logProps.put(AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG, tableTopicSchemaRegistryUrl) + } + // AutoMQ inject end + logProps } diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala index 5976653cf0..8b38ae1c9e 100755 --- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala +++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala @@ -216,7 +216,7 @@ class KafkaRequestHandlerPool( this.logIdent = "[" + logAndThreadNamePrefix + " Kafka Request Handler on Broker " + brokerId + "], " val runnables = new mutable.ArrayBuffer[KafkaRequestHandler](numThreads) - var multiRequestQueue = requestChannel.registerNRequestHandler(numThreads) + requestChannel.registerNRequestHandler(numThreads) for (i <- 0 until numThreads) { createHandler(i) @@ -541,10 +541,6 @@ class BrokerTopicStats(remoteStorageEnabled: Boolean = false) extends Logging { def topicStats(topic: String): BrokerTopicMetrics = stats.getAndMaybePut(topic) - def topicPartitionStats(topicPartition: TopicPartition): BrokerTopicPartitionMetrics = { - partitionStats.getAndMaybePut(topicPartition) - } - def updateReplicationBytesIn(value: Long): Unit = { allTopicsStats.replicationBytesInRate.foreach { metric => metric.mark(value) @@ -649,12 +645,32 @@ class BrokerTopicStats(remoteStorageEnabled: Boolean = false) extends Logging { updateReassignmentBytesOut(value) updateReplicationBytesOut(value) } else { - topicPartitionStats(topicPartition).bytesOutRate.mark(value) + // AutoMQ inject start + val partitionMetrics = partitionStats.getAndMaybePut(topicPartition) + partitionMetrics.bytesOutRate.mark(value) + // AutoMQ inject start topicStats(topicPartition.topic()).bytesOutRate.mark(value) allTopicsStats.bytesOutRate.mark(value) } } + // AutoMQ inject start + def newMetrics(topicPartition: TopicPartition): Unit = { + partitionStats.getAndMaybePut(topicPartition) + } + + def updatePartitionBytesIn(topicPartition: TopicPartition, size: Int): Unit = { + val partitionMetrics = partitionStats.getAndMaybePut(topicPartition) + partitionMetrics.totalProduceRequestRate.mark() + partitionMetrics.bytesInRate.mark(size) + } + + def updatePartitionFetchRequestRate(topicPartition: TopicPartition): Unit = { + val partitionMetrics = partitionStats.getAndMaybePut(topicPartition) + partitionMetrics.totalFetchRequestRate.mark() + } + // AutoMQ inject end + // Update the broker-level all topic metric values so that we have a sample right for all topics metric after update of partition. def recordRemoteCopyLagBytes(topic: String, partition: Int, value: Long): Unit = { val topicMetric = topicStats(topic) diff --git a/core/src/main/scala/kafka/server/Limiter.java b/core/src/main/scala/kafka/server/Limiter.java index 3e3511de77..0fcb3505c2 100644 --- a/core/src/main/scala/kafka/server/Limiter.java +++ b/core/src/main/scala/kafka/server/Limiter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.server; @@ -49,6 +57,16 @@ public interface Limiter { */ int availablePermits(); + /** + * Return the number of threads waiting for permits. + */ + int waitingThreads(); + + /** + * Return the name of this limiter. + */ + String name(); + /** * A handler to release acquired permits. */ diff --git a/core/src/main/scala/kafka/server/MetadataCache.scala b/core/src/main/scala/kafka/server/MetadataCache.scala index e38ffe0c92..1d9e0c43d9 100755 --- a/core/src/main/scala/kafka/server/MetadataCache.scala +++ b/core/src/main/scala/kafka/server/MetadataCache.scala @@ -22,11 +22,14 @@ import org.apache.kafka.admin.BrokerMetadata import org.apache.kafka.common.message.{MetadataResponseData, UpdateMetadataRequestData} import org.apache.kafka.common.network.ListenerName import org.apache.kafka.common.{Cluster, Node, TopicPartition, Uuid} +import org.apache.kafka.image.MetadataImage import org.apache.kafka.metadata.BrokerRegistration import org.apache.kafka.server.common.automq.AutoMQVersion import org.apache.kafka.server.common.{FinalizedFeatures, KRaftVersion, MetadataVersion} +import java.nio.ByteBuffer import java.util +import java.util.OptionalLong import java.util.function.Supplier import scala.collection._ @@ -120,6 +123,12 @@ trait MetadataCache { def getPartitionLeaderNode(topicName: String, partitionId: Int): BrokerRegistration def getNode(nodeId: Int): BrokerRegistration + + def getValue(key: String): ByteBuffer + + def getStreamEndOffset(streamId: Long): OptionalLong + + def safeRun[T](func: Function[MetadataImage, T]): T // AutoMQ inject end } diff --git a/core/src/main/scala/kafka/server/NoopLimiter.java b/core/src/main/scala/kafka/server/NoopLimiter.java index 1fab234492..d4ea50cd4b 100644 --- a/core/src/main/scala/kafka/server/NoopLimiter.java +++ b/core/src/main/scala/kafka/server/NoopLimiter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.server; @@ -38,6 +46,16 @@ public int availablePermits() { return Integer.MAX_VALUE; } + @Override + public int waitingThreads() { + return 0; + } + + @Override + public String name() { + return "noop"; + } + public static class NoopHandler implements Handler { @Override public void close() { diff --git a/core/src/main/scala/kafka/server/QuotaFactory.scala b/core/src/main/scala/kafka/server/QuotaFactory.scala index 990565e04c..9a13d864f7 100644 --- a/core/src/main/scala/kafka/server/QuotaFactory.scala +++ b/core/src/main/scala/kafka/server/QuotaFactory.scala @@ -32,6 +32,16 @@ object QuotaType { case object Fetch extends QuotaType case object Produce extends QuotaType case object Request extends QuotaType + // AutoMQ for Kafka inject start + /** + * Quota type for slow fetch throughput limiting. + */ + case object SlowFetch extends QuotaType + /** + * Quota type for request rate limiting. + */ + case object RequestRate extends QuotaType + // AutoMQ for Kafka inject end case object ControllerMutation extends QuotaType case object LeaderReplication extends QuotaType case object FollowerReplication extends QuotaType @@ -44,11 +54,16 @@ object QuotaType { case QuotaType.Fetch => ClientQuotaType.FETCH case QuotaType.Produce => ClientQuotaType.PRODUCE case QuotaType.Request => ClientQuotaType.REQUEST + // AutoMQ for Kafka inject start + case QuotaType.SlowFetch => ClientQuotaType.SLOW_FETCH + case QuotaType.RequestRate => ClientQuotaType.REQUEST_RATE + // AutoMQ for Kafka inject end case QuotaType.ControllerMutation => ClientQuotaType.CONTROLLER_MUTATION case _ => throw new IllegalArgumentException(s"Not a client quota type: $quotaType") } } + // AutoMQ for Kafka inject start // for test def fetch(): QuotaType = { QuotaType.Fetch @@ -58,9 +73,14 @@ object QuotaType { QuotaType.Produce } - def request(): QuotaType = { - QuotaType.Request + def slowFetch(): QuotaType = { + QuotaType.SlowFetch + } + + def requestRate(): QuotaType = { + QuotaType.RequestRate } + // AutoMQ for Kafka inject end } sealed trait QuotaType diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala index ee028d64e1..e61867b775 100644 --- a/core/src/main/scala/kafka/server/ReplicaManager.scala +++ b/core/src/main/scala/kafka/server/ReplicaManager.scala @@ -67,7 +67,7 @@ import java.nio.file.{Files, Paths} import java.util import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong} import java.util.concurrent.locks.Lock -import java.util.concurrent.{ArrayBlockingQueue, CompletableFuture, Future, RejectedExecutionException, TimeUnit} +import java.util.concurrent.{CompletableFuture, Future, LinkedBlockingQueue, RejectedExecutionException, TimeUnit} import java.util.{Collections, Optional, OptionalInt, OptionalLong} import scala.collection.{Map, Seq, Set, mutable} import scala.compat.java8.OptionConverters._ @@ -939,7 +939,8 @@ class ReplicaManager(val config: KafkaConfig, case class Verification( hasInflight: AtomicBoolean, - waitingRequests: ArrayBlockingQueue[TransactionVerificationRequest], + // Use an unbounded queue to prevent deadlock from happening. See https://github.com/AutoMQ/automq/issues/2902 + waitingRequests: LinkedBlockingQueue[TransactionVerificationRequest], timestamp: AtomicLong, ) diff --git a/core/src/main/scala/kafka/server/SharedServer.scala b/core/src/main/scala/kafka/server/SharedServer.scala index cf2d8b3fd0..0f9b1ed462 100644 --- a/core/src/main/scala/kafka/server/SharedServer.scala +++ b/core/src/main/scala/kafka/server/SharedServer.scala @@ -17,7 +17,7 @@ package kafka.server -import kafka.log.stream.s3.telemetry.TelemetryManager +import com.automq.opentelemetry.AutoMQTelemetryManager import kafka.raft.KafkaRaftManager import kafka.server.Server.MetricsPrefix import kafka.server.metadata.BrokerServerMetrics @@ -32,9 +32,8 @@ import org.apache.kafka.image.loader.MetadataLoader import org.apache.kafka.image.loader.metrics.MetadataLoaderMetrics import org.apache.kafka.image.publisher.metrics.SnapshotEmitterMetrics import org.apache.kafka.image.publisher.{SnapshotEmitter, SnapshotGenerator} -import org.apache.kafka.metadata.ListenerInfo -import org.apache.kafka.metadata.MetadataRecordSerde import org.apache.kafka.metadata.properties.MetaPropertiesEnsemble +import org.apache.kafka.metadata.{ListenerInfo, MetadataRecordSerde} import org.apache.kafka.raft.Endpoints import org.apache.kafka.server.ProcessRole import org.apache.kafka.server.common.ApiMessageAndVersion @@ -42,12 +41,9 @@ import org.apache.kafka.server.fault.{FaultHandler, LoggingFaultHandler, Process import org.apache.kafka.server.metrics.KafkaYammerMetrics import java.net.InetSocketAddress -import java.util.Arrays -import java.util.Optional import java.util.concurrent.atomic.AtomicReference import java.util.concurrent.{CompletableFuture, TimeUnit} -import java.util.{Collection => JCollection} -import java.util.{Map => JMap} +import java.util.{Arrays, Optional, Collection => JCollection, Map => JMap} import scala.jdk.CollectionConverters._ /** @@ -113,7 +109,7 @@ class SharedServer( // AutoMQ for Kafka injection start ElasticStreamSwitch.setSwitch(sharedServerConfig.elasticStreamEnabled) - @volatile var telemetryManager: TelemetryManager = _ + @volatile var telemetryManager: AutoMQTelemetryManager = _ // AutoMQ for Kafka injection end @volatile var metrics: Metrics = _metrics @@ -130,10 +126,10 @@ class SharedServer( def nodeId: Int = metaPropsEnsemble.nodeId().getAsInt - protected def buildTelemetryManager(config: KafkaConfig, clusterId: String): TelemetryManager = { - new TelemetryManager(config, clusterId) + protected def buildTelemetryManager(config: KafkaConfig, clusterId: String): AutoMQTelemetryManager = { + TelemetrySupport.start(config, clusterId) } - + private def isUsed(): Boolean = synchronized { usedByController || usedByBroker } @@ -287,7 +283,6 @@ class SharedServer( // AutoMQ inject start telemetryManager = buildTelemetryManager(sharedServerConfig, clusterId) - telemetryManager.init() // AutoMQ inject end val _raftManager = new KafkaRaftManager[ApiMessageAndVersion]( diff --git a/core/src/main/scala/kafka/server/metadata/BrokerMetadataPublisher.scala b/core/src/main/scala/kafka/server/metadata/BrokerMetadataPublisher.scala index 627744900a..559fd54eff 100644 --- a/core/src/main/scala/kafka/server/metadata/BrokerMetadataPublisher.scala +++ b/core/src/main/scala/kafka/server/metadata/BrokerMetadataPublisher.scala @@ -33,7 +33,6 @@ import org.apache.kafka.image.publisher.MetadataPublisher import org.apache.kafka.image.{MetadataDelta, MetadataImage, TopicDelta, TopicsDelta} import org.apache.kafka.server.fault.FaultHandler -import java.util import java.util.concurrent.{CompletableFuture, ExecutorService} import scala.collection.mutable import scala.jdk.CollectionConverters._ @@ -343,7 +342,7 @@ class BrokerMetadataPublisher( ) } catch { case t: Throwable => metadataPublishingFaultHandler.handleFault("Error updating group " + - s"coordinator with local changes in $deltaName", t) + s"coordinator with local changes in $deltaName for $topicPartition", t) } } @@ -357,22 +356,10 @@ class BrokerMetadataPublisher( txnCoordinator.onResignation) } catch { case t: Throwable => metadataPublishingFaultHandler.handleFault("Error updating txn " + - s"coordinator with local changes in $deltaName", t) + s"coordinator with local changes in $deltaName for $topicPartition", t) } } }) - - try { - // Notify the group coordinator about deleted topics. - if (topicsDelta.topicWasDeleted(topicPartition.topic())) { - groupCoordinator.onPartitionsDeleted( - util.List.of(new TopicPartition(topicPartition.topic(), topicPartition.partition())), - RequestLocal.NoCaching.bufferSupplier) - } - } catch { - case t: Throwable => metadataPublishingFaultHandler.handleFault("Error updating group " + - s"coordinator with deleted partitions in $deltaName", t) - } } // Notify the replica manager about changes to topics. @@ -383,6 +370,7 @@ class BrokerMetadataPublisher( if (ex != null) { metadataPublishingFaultHandler.handleFault("Error applying topics " + s"delta in $deltaName", ex) } + notifyGroupCoordinatorOfDeletedPartitions(deltaName, topicsDelta) }) case _ => CompletableFuture.completedFuture(replicaManager.applyDelta(topicsDelta, newImage)) @@ -413,25 +401,29 @@ class BrokerMetadataPublisher( case t: Throwable => metadataPublishingFaultHandler.handleFault("Error updating txn " + s"coordinator with local changes in $deltaName", t) } - try { - // Notify the group coordinator about deleted topics. - val deletedTopicPartitions = new mutable.ArrayBuffer[TopicPartition]() - topicsDelta.deletedTopicIds().forEach { id => - val topicImage = topicsDelta.image().getTopic(id) - topicImage.partitions().keySet().forEach { - id => deletedTopicPartitions += new TopicPartition(topicImage.name(), id) - } - } - if (deletedTopicPartitions.nonEmpty) { - groupCoordinator.onPartitionsDeleted(deletedTopicPartitions.asJava, RequestLocal.NoCaching.bufferSupplier) - } - } catch { - case t: Throwable => metadataPublishingFaultHandler.handleFault("Error updating group " + - s"coordinator with deleted partitions in $deltaName", t) - } + notifyGroupCoordinatorOfDeletedPartitions(deltaName, topicsDelta) }) } } + + private def notifyGroupCoordinatorOfDeletedPartitions(deltaName: String, topicsDelta: TopicsDelta): Unit = { + try { + // Notify the group coordinator about deleted topics. + val deletedTopicPartitions = new mutable.ArrayBuffer[TopicPartition]() + topicsDelta.deletedTopicIds().forEach { id => + val topicImage = topicsDelta.image().getTopic(id) + topicImage.partitions().keySet().forEach { + id => deletedTopicPartitions += new TopicPartition(topicImage.name(), id) + } + } + if (deletedTopicPartitions.nonEmpty) { + groupCoordinator.onPartitionsDeleted(deletedTopicPartitions.asJava, RequestLocal.NoCaching.bufferSupplier) + } + } catch { + case t: Throwable => metadataPublishingFaultHandler.handleFault("Error updating group " + + s"coordinator with deleted partitions in $deltaName", t) + } + } // AutoMQ inject end } diff --git a/core/src/main/scala/kafka/server/metadata/KRaftMetadataCache.scala b/core/src/main/scala/kafka/server/metadata/KRaftMetadataCache.scala index 1e25cb7592..1ac8aa3b6a 100644 --- a/core/src/main/scala/kafka/server/metadata/KRaftMetadataCache.scala +++ b/core/src/main/scala/kafka/server/metadata/KRaftMetadataCache.scala @@ -40,11 +40,12 @@ import org.apache.kafka.metadata.{BrokerRegistration, PartitionRegistration, Rep import org.apache.kafka.server.common.automq.AutoMQVersion import org.apache.kafka.server.common.{FinalizedFeatures, KRaftVersion, MetadataVersion} +import java.nio.ByteBuffer import java.util import java.util.concurrent.ThreadLocalRandom import java.util.concurrent.locks.ReentrantLock import java.util.function.Supplier -import java.util.{Collections, Properties} +import java.util.{Collections, OptionalLong, Properties} import scala.collection.mutable.ListBuffer import scala.collection.{Map, Seq, Set, mutable} import scala.compat.java8.OptionConverters._ @@ -118,7 +119,7 @@ class KRaftMetadataCache( streamMetadata.state == StreamState.OPENED && streamMetadata.lastRange().nodeId() == tpRegistration.leader - if (!result) { + if (!result && isDebugEnabled) { debug(s"Failover failed for topicPartition $topicPartition, tpEpoch $tpRegistration, streamMetadata ${streamMetadata}") } @@ -622,7 +623,7 @@ class KRaftMetadataCache( _currentImage.features().autoMQVersion() } - def safeRun[T](func: Function[MetadataImage, T]): T = { + override def safeRun[T](func: Function[MetadataImage, T]): T = { val image = retainedImage() try { func.apply(image) @@ -631,7 +632,7 @@ class KRaftMetadataCache( } } - private def retainedImage(): MetadataImage = { + def retainedImage(): MetadataImage = { imageLock.lock() try { val image = _currentImage @@ -658,7 +659,14 @@ class KRaftMetadataCache( override def getNode(nodeId: Int): BrokerRegistration = { _currentImage.cluster().broker(nodeId) } + + override def getValue(key: String): ByteBuffer = { + _currentImage.kv().getValue(key) + } + + override def getStreamEndOffset(streamId: Long): OptionalLong = { + _currentImage.streamsMetadata().streamEndOffset(streamId) + } // AutoMQ inject end } - diff --git a/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala b/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala index 2052109c18..731372ed9d 100755 --- a/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala +++ b/core/src/main/scala/kafka/server/metadata/ZkMetadataCache.scala @@ -18,7 +18,7 @@ package kafka.server.metadata import java.util -import java.util.{Collections, Optional} +import java.util.{Collections, Optional, OptionalLong} import java.util.concurrent.locks.{ReentrantLock, ReentrantReadWriteLock} import scala.collection.{Seq, Set, mutable} import scala.jdk.CollectionConverters._ @@ -40,10 +40,12 @@ import org.apache.kafka.common.network.ListenerName import org.apache.kafka.common.protocol.Errors import org.apache.kafka.common.requests.{AbstractControlRequest, ApiVersionsResponse, MetadataResponse, UpdateMetadataRequest} import org.apache.kafka.common.security.auth.SecurityProtocol +import org.apache.kafka.image.MetadataImage import org.apache.kafka.metadata.BrokerRegistration import org.apache.kafka.server.common.automq.AutoMQVersion import org.apache.kafka.server.common.{FinalizedFeatures, MetadataVersion} +import java.nio.ByteBuffer import java.util.concurrent.{ThreadLocalRandom, TimeUnit} import scala.concurrent.TimeoutException import scala.math.max @@ -727,6 +729,18 @@ class ZkMetadataCache( override def getNode(nodeId: Int): BrokerRegistration = { throw new UnsupportedOperationException() } + + override def getValue(key: String): ByteBuffer = { + throw new UnsupportedOperationException() + } + + override def getStreamEndOffset(streamId: Long): OptionalLong = { + throw new UnsupportedOperationException() + } + + override def safeRun[T](func: Function[MetadataImage, T]): T = { + throw new UnsupportedOperationException() + } // AutoMQ inject end } diff --git a/core/src/main/scala/kafka/server/streamaspect/BrokerQuotaManager.scala b/core/src/main/scala/kafka/server/streamaspect/BrokerQuotaManager.scala index accfe469be..a206a5e7db 100644 --- a/core/src/main/scala/kafka/server/streamaspect/BrokerQuotaManager.scala +++ b/core/src/main/scala/kafka/server/streamaspect/BrokerQuotaManager.scala @@ -1,29 +1,38 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.server.streamaspect +import com.automq.stream.s3.metrics.S3StreamMetricsManager import kafka.network.RequestChannel import kafka.server._ import kafka.utils.QuotaUtils import org.apache.kafka.common.MetricName import org.apache.kafka.common.metrics.stats.{Avg, CumulativeSum, Rate} import org.apache.kafka.common.metrics.{Metrics, Quota, QuotaViolationException, Sensor} -import org.apache.kafka.common.requests.RequestContext import org.apache.kafka.common.security.auth.KafkaPrincipal import org.apache.kafka.common.utils.Time import org.apache.kafka.network.Session -import org.apache.kafka.server.config.BrokerQuotaManagerConfig +import org.apache.kafka.server.config.{BrokerQuotaManagerConfig, QuotaConfigs} -import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.{Optional, Properties} import scala.collection.mutable import scala.jdk.CollectionConverters._ @@ -32,6 +41,7 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, private val time: Time, private val threadNamePrefix: String) extends ClientRequestQuotaManager(config, metrics, time, threadNamePrefix, None) { + private val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(this.config.quotaWindowSizeSeconds * this.config.numQuotaSamples) private val metricsTags = Map("domain" -> "broker", "nodeId" -> String.valueOf(config.nodeId())) private val whiteListCache = mutable.HashMap[String, Boolean]() @@ -41,14 +51,31 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, override def delayQueueSensor: Sensor = brokerDelayQueueSensor - def getMaxValueInQuotaWindow(quotaType: QuotaType): Double = { - if (config.quotaEnabled) { + S3StreamMetricsManager.registerBrokerQuotaLimitSupplier(() => java.util.Map.of( + QuotaType.RequestRate.toString, quotaLimit(QuotaType.RequestRate), + QuotaType.Produce.toString, quotaLimit(QuotaType.Produce), + QuotaType.Fetch.toString, quotaLimit(QuotaType.Fetch), + QuotaType.SlowFetch.toString, quotaLimit(QuotaType.SlowFetch) + )) + + def getMaxValueInQuotaWindow(quotaType: QuotaType, request: RequestChannel.Request): Double = { + if (shouldThrottle(request)) { quotaLimit(quotaType) } else { Double.MaxValue } } + /** + * Get the value of the metric for the given quota type at the given time. + * It return empty if the metric is not found, which is possible if the quota is disabled or no request has been + * processed yet. + */ + def getQuotaMetricValue(quotaType: QuotaType, timeMs: Long): Optional[java.lang.Double] = { + Optional.ofNullable(metrics.metric(clientQuotaMetricName(quotaType, metricsTags))) + .map(_.measurableValueV2(timeMs)) + } + def recordNoThrottle(quotaType: QuotaType, value: Double): Unit = { val clientSensors = getOrCreateQuotaSensors(quotaType) clientSensors.quotaSensor.record(value, time.milliseconds(), false) @@ -56,19 +83,26 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, def maybeRecordAndGetThrottleTimeMs(quotaType: QuotaType, request: RequestChannel.Request, value: Double, timeMs: Long): Int = { - if (!config.quotaEnabled) { - return 0 + if (shouldThrottle(request)) { + maybeRecordAndGetThrottleTimeMs(quotaType, value, timeMs) + } else { + 0 } - - maybeRecordAndGetThrottleTimeMs(quotaType, request.session, request.context, value, timeMs) } - protected def throttleTime(quotaType: QuotaType, e: QuotaViolationException, timeMs: Long): Long = { - if (quotaType == QuotaType.Request) { + override protected def throttleTime(e: QuotaViolationException, timeMs: Long): Long = { QuotaUtils.boundedThrottleTime(e, maxThrottleTimeMs, timeMs) - } else { - QuotaUtils.throttleTime(e, timeMs) - } + } + + private def shouldThrottle(request: RequestChannel.Request): Boolean = { + val quotaEnabled = config.quotaEnabled + val isInternal = isInternalClient(request.context.clientId()) + val isWhiteListed = isInWhiteList(request.session.principal, request.context.clientId(), request.context.listenerName()) + quotaEnabled && !isInternal && !isWhiteListed + } + + private def isInternalClient(clientId: String): Boolean = { + clientId.startsWith(QuotaConfigs.INTERNAL_CLIENT_ID_PREFIX) } private def isInWhiteList(principal: KafkaPrincipal, clientId: String, listenerName: String): Boolean = { @@ -84,19 +118,17 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, } } - def maybeRecordAndGetThrottleTimeMs(quotaType: QuotaType, session: Session, context: RequestContext, value: Double, - timeMs: Long): Int = { - if (isInWhiteList(session.principal, context.clientId(), context.listenerName())) { - return 0 - } + def maybeRecordAndGetThrottleTimeMs(quotaType: QuotaType, value: Double, timeMs: Long): Int = { val clientSensors = getOrCreateQuotaSensors(quotaType) try { clientSensors.quotaSensor.record(value, timeMs, true) 0 } catch { case e: QuotaViolationException => - val throttleTimeMs = throttleTime(quotaType, e, timeMs).toInt - debug(s"Quota violated for sensor (${clientSensors.quotaSensor.name}). Delay time: ($throttleTimeMs)") + val throttleTimeMs = throttleTime(e, timeMs).toInt + if (isDebugEnabled) { + debug(s"Quota violated for sensor (${clientSensors.quotaSensor.name}). Delay time: ($throttleTimeMs)") + } throttleTimeMs } } @@ -112,31 +144,56 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, whiteListCache.clear() if (!config.quotaEnabled) { - metrics.removeSensor(getQuotaSensorName(QuotaType.Request, metricsTags)) + metrics.removeSensor(getQuotaSensorName(QuotaType.RequestRate, metricsTags)) metrics.removeSensor(getQuotaSensorName(QuotaType.Produce, metricsTags)) metrics.removeSensor(getQuotaSensorName(QuotaType.Fetch, metricsTags)) - metrics.removeSensor(getThrottleTimeSensorName(QuotaType.Request, metricsTags)) + metrics.removeSensor(getQuotaSensorName(QuotaType.SlowFetch, metricsTags)) + metrics.removeSensor(getThrottleTimeSensorName(QuotaType.RequestRate, metricsTags)) metrics.removeSensor(getThrottleTimeSensorName(QuotaType.Produce, metricsTags)) metrics.removeSensor(getThrottleTimeSensorName(QuotaType.Fetch, metricsTags)) + metrics.removeSensor(getThrottleTimeSensorName(QuotaType.SlowFetch, metricsTags)) return } val allMetrics = metrics.metrics() - val requestMetrics = allMetrics.get(clientQuotaMetricName(QuotaType.Request, metricsTags)) - if (requestMetrics != null) { - requestMetrics.config(getQuotaMetricConfig(quotaLimit(QuotaType.Request))) + val requestRateMetric = allMetrics.get(clientQuotaMetricName(QuotaType.RequestRate, metricsTags)) + if (requestRateMetric != null) { + requestRateMetric.config(getQuotaMetricConfig(quotaLimit(QuotaType.RequestRate))) } - val produceMetrics = allMetrics.get(clientQuotaMetricName(QuotaType.Produce, metricsTags)) - if (produceMetrics != null) { - produceMetrics.config(getQuotaMetricConfig(quotaLimit(QuotaType.Produce))) + val produceMetric = allMetrics.get(clientQuotaMetricName(QuotaType.Produce, metricsTags)) + if (produceMetric != null) { + produceMetric.config(getQuotaMetricConfig(quotaLimit(QuotaType.Produce))) } - val fetchMetrics = allMetrics.get(clientQuotaMetricName(QuotaType.Fetch, metricsTags)) - if (fetchMetrics != null) { - fetchMetrics.config(getQuotaMetricConfig(quotaLimit(QuotaType.Fetch))) + val fetchMetric = allMetrics.get(clientQuotaMetricName(QuotaType.Fetch, metricsTags)) + if (fetchMetric != null) { + fetchMetric.config(getQuotaMetricConfig(quotaLimit(QuotaType.Fetch))) } + + val slowFetchMetric = allMetrics.get(clientQuotaMetricName(QuotaType.SlowFetch, metricsTags)) + if (slowFetchMetric != null) { + slowFetchMetric.config(getQuotaMetricConfig(quotaLimit(QuotaType.SlowFetch))) + } + } + } + + def updateQuota(quotaType: QuotaType, quota: Double): Unit = { + // update the quota in the config first to make sure the new quota will be used if {@link #updateQuotaMetricConfigs} is called + quotaType match { + case QuotaType.RequestRate => config.requestRateQuota(quota) + case QuotaType.Produce => config.produceQuota(quota) + case QuotaType.Fetch => config.fetchQuota(quota) + case QuotaType.SlowFetch => config.slowFetchQuota(quota) + case _ => throw new IllegalArgumentException(s"Unknown quota type $quotaType") + } + + // update the metric config + val allMetrics = metrics.metrics() + val metric = allMetrics.get(clientQuotaMetricName(quotaType, metricsTags)) + if (metric != null) { + metric.config(getQuotaMetricConfig(quotaLimit(quotaType))) } } @@ -151,7 +208,9 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, val throttledChannel = new ThrottledChannel(time, throttleTimeMs, throttleCallback) delayQueue.add(throttledChannel) delayQueueSensor.record() - debug("Channel throttled for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs)) + if (isDebugEnabled) { + debug("Channel throttled for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs)) + } } } @@ -161,11 +220,14 @@ class BrokerQuotaManager(private val config: BrokerQuotaManagerConfig, private def getQuotaSensorName(quotaType: QuotaType, metricTags: Map[String, String]): String = s"$quotaType-${metricTagsToSensorSuffix(metricTags)}" - private def quotaLimit(quotaType: QuotaType): Double = { - if (quotaType == QuotaType.Request) config.requestQuota - else if (quotaType == QuotaType.Produce) config.produceQuota - else if (quotaType == QuotaType.Fetch) config.fetchQuota - else throw new IllegalArgumentException(s"Unknown quota type $quotaType") + def quotaLimit(quotaType: QuotaType): Double = { + quotaType match { + case QuotaType.RequestRate => config.requestRateQuota + case QuotaType.Produce => config.produceQuota + case QuotaType.Fetch => config.fetchQuota + case QuotaType.SlowFetch => config.slowFetchQuota + case _ => throw new IllegalArgumentException(s"Unknown quota type $quotaType") + } } protected def clientQuotaMetricName(quotaType: QuotaType, quotaMetricTags: Map[String, String]): MetricName = { diff --git a/core/src/main/scala/kafka/server/streamaspect/ElasticKafkaApis.scala b/core/src/main/scala/kafka/server/streamaspect/ElasticKafkaApis.scala index 40f02ad9bd..fb0712ce4f 100644 --- a/core/src/main/scala/kafka/server/streamaspect/ElasticKafkaApis.scala +++ b/core/src/main/scala/kafka/server/streamaspect/ElasticKafkaApis.scala @@ -1,10 +1,13 @@ package kafka.server.streamaspect import com.automq.stream.s3.metrics.TimerUtil +import com.automq.stream.s3.network.{GlobalNetworkBandwidthLimiters, ThrottleStrategy} +import com.automq.stream.utils.Threads +import com.automq.stream.utils.threads.S3StreamThreadPoolMonitor import com.yammer.metrics.core.Histogram -import kafka.automq.zonerouter.{ClientIdMetadata, NoopProduceRouter, ProduceRouter} +import kafka.automq.interceptor.{ClientIdMetadata, NoopTrafficInterceptor, ProduceRequestArgs, TrafficInterceptor} import kafka.coordinator.transaction.TransactionCoordinator -import kafka.log.streamaspect.ElasticLogManager +import kafka.log.streamaspect.{ElasticLogManager, ReadHint} import kafka.metrics.KafkaMetricsUtil import kafka.network.RequestChannel import kafka.server.QuotaFactory.QuotaManagers @@ -13,7 +16,6 @@ import kafka.server.metadata.ConfigRepository import kafka.server.streamaspect.ElasticKafkaApis.{LAST_RECORD_TIMESTAMP, PRODUCE_ACK_TIME_HIST, PRODUCE_CALLBACK_TIME_HIST, PRODUCE_TIME_HIST} import kafka.utils.Implicits.MapExtensionMethods import org.apache.kafka.admin.AdminUtils -import org.apache.kafka.common.{Node, TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.common.acl.AclOperation.{CLUSTER_ACTION, READ, WRITE} import org.apache.kafka.common.errors.{ApiException, UnsupportedCompressionTypeException} import org.apache.kafka.common.internals.FatalExitError @@ -21,27 +23,30 @@ import org.apache.kafka.common.message.{DeleteTopicsRequestData, FetchResponseDa import org.apache.kafka.common.metrics.Metrics import org.apache.kafka.common.network.{ListenerName, NetworkSend, Send} import org.apache.kafka.common.protocol.{ApiKeys, Errors} -import org.apache.kafka.common.record.{LazyDownConversionRecords, MemoryRecords, MultiRecordsSend, PooledResource, RecordBatch, RecordValidationStats} +import org.apache.kafka.common.record._ import org.apache.kafka.common.replica.ClientMetadata import org.apache.kafka.common.replica.ClientMetadata.DefaultClientMetadata import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse -import org.apache.kafka.common.requests.s3.AutomqZoneRouterRequest +import org.apache.kafka.common.requests.s3.{AutomqGetPartitionSnapshotRequest, AutomqUpdateGroupRequest, AutomqUpdateGroupResponse, AutomqZoneRouterRequest} import org.apache.kafka.common.requests.{AbstractResponse, DeleteTopicsRequest, DeleteTopicsResponse, FetchRequest, FetchResponse, ProduceRequest, ProduceResponse, RequestUtils} import org.apache.kafka.common.resource.Resource.CLUSTER_NAME import org.apache.kafka.common.resource.ResourceType.{CLUSTER, TOPIC, TRANSACTIONAL_ID} -import org.apache.kafka.common.utils.{ThreadUtils, Time} +import org.apache.kafka.common.utils.Time +import org.apache.kafka.common.{Node, TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.coordinator.group.GroupCoordinator import org.apache.kafka.server.ClientMetricsManager import org.apache.kafka.server.authorizer.Authorizer import org.apache.kafka.server.metrics.KafkaMetricsGroup import org.apache.kafka.server.record.BrokerCompressionType import org.apache.kafka.storage.internals.log.{FetchIsolation, FetchParams, FetchPartitionData} +import org.slf4j.LoggerFactory import java.util -import java.util.{Collections, Optional} import java.util.concurrent.atomic.AtomicLong -import java.util.concurrent.{ExecutorService, Executors, TimeUnit} +import java.util.concurrent.{CompletableFuture, ExecutorService, TimeUnit} +import java.util.function.Supplier import java.util.stream.IntStream +import java.util.{Collections, Optional} import scala.annotation.nowarn import scala.collection.{Map, Seq, mutable} import scala.jdk.CollectionConverters.{CollectionHasAsScala, MapHasAsJava, MapHasAsScala, SeqHasAsJava} @@ -76,13 +81,16 @@ class ElasticKafkaApis( tokenManager: DelegationTokenManager, apiVersionManager: ApiVersionManager, clientMetricsManager: Option[ClientMetricsManager], - val asyncHandleExecutor: ExecutorService = Executors.newSingleThreadExecutor(ThreadUtils.createThreadFactory("kafka-apis-async-handle-executor-%d", true)), - val listOffsetHandleExecutor: ExecutorService = Executors.newSingleThreadExecutor(ThreadUtils.createThreadFactory("kafka-apis-list-offset-handle-executor-%d", true)) + val deleteTopicHandleExecutor: ExecutorService = S3StreamThreadPoolMonitor.createAndMonitor(1, 1, 0L, TimeUnit.MILLISECONDS, "kafka-apis-delete-topic-handle-executor", true, 1000), + val listOffsetHandleExecutor: ExecutorService = S3StreamThreadPoolMonitor.createAndMonitor(1, 1, 0L, TimeUnit.MILLISECONDS, "kafka-apis-list-offset-handle-executor", true, 1000) ) extends KafkaApis(requestChannel, metadataSupport, replicaManager, groupCoordinator, txnCoordinator, autoTopicCreationManager, brokerId, config, configRepository, metadataCache, metrics, authorizer, quotas, fetchManager, brokerTopicStats, clusterId, time, tokenManager, apiVersionManager, clientMetricsManager) { - private var produceRouter: ProduceRouter = new NoopProduceRouter(this, metadataCache) + private val offsetForLeaderEpochExecutor: ExecutorService = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(1, "kafka-apis-offset-for-leader-epoch-handle-executor", true, LoggerFactory.getLogger(ElasticKafkaApis.getClass)) + + private var trafficInterceptor: TrafficInterceptor = new NoopTrafficInterceptor(this, metadataCache) + private var snapshotAwaitReadySupplier: Supplier[CompletableFuture[Void]] = () => CompletableFuture.completedFuture(null) /** * Generate a map of topic -> [(partitionId, epochId)] based on provided topicsRequestData. @@ -131,7 +139,7 @@ class ElasticKafkaApis( response.asInstanceOf[DeleteTopicsResponse].data().responses().forEach(result => { if (result.errorCode() == Errors.NONE.code()) { if (!metadataCache.autoMQVersion().isTopicCleanupByControllerSupported) { - asyncHandleExecutor.submit(new Runnable { + deleteTopicHandleExecutor.submit(new Runnable { override def run(): Unit = { topicNameToPartitionEpochsMap.get(result.name()).foreach(partitionEpochs => { ElasticLogManager.destroyLog(new TopicPartition(result.name(), partitionEpochs._1), result.topicId(), partitionEpochs._2) @@ -167,8 +175,11 @@ class ElasticKafkaApis( request.header.apiKey match { case ApiKeys.AUTOMQ_ZONE_ROUTER => handleZoneRouterRequest(request, requestLocal) + case ApiKeys.AUTOMQ_GET_PARTITION_SNAPSHOT => handleGetPartitionSnapshotRequest(request, requestLocal) case ApiKeys.DELETE_TOPICS => maybeForwardTopicDeletionToController(request, handleDeleteTopicsRequest) case ApiKeys.GET_NEXT_NODE_ID => forwardToControllerOrFail(request) + case ApiKeys.AUTOMQ_UPDATE_GROUP => handleUpdateGroupRequest(request, requestLocal) + case _ => throw new IllegalStateException("Message conversion info is recorded only for Produce/Fetch requests") } @@ -193,12 +204,16 @@ class ElasticKafkaApis( override def handle(request: RequestChannel.Request, requestLocal: RequestLocal): Unit = { request.header.apiKey match { - case ApiKeys.DELETE_TOPICS | ApiKeys.GET_NEXT_NODE_ID | ApiKeys.AUTOMQ_ZONE_ROUTER => handleExtensionRequest(request, requestLocal) + case ApiKeys.DELETE_TOPICS + | ApiKeys.GET_NEXT_NODE_ID + | ApiKeys.AUTOMQ_ZONE_ROUTER + | ApiKeys.AUTOMQ_UPDATE_GROUP + | ApiKeys.AUTOMQ_GET_PARTITION_SNAPSHOT => handleExtensionRequest(request, requestLocal) case _ => super.handle(request, requestLocal) } } - protected def getCurrentLeaderForProduce(tp: TopicPartition, clientIdMetadata: ClientIdMetadata, ln: ListenerName): LeaderNode = { + protected def getCurrentLeaderForProduceAndFetch(tp: TopicPartition, ln: ListenerName, clientIdMetadata: ClientIdMetadata): LeaderNode = { val partitionInfoOrError = replicaManager.getPartitionOrError(tp) val (leaderId, leaderEpoch) = partitionInfoOrError match { case Right(x) => @@ -210,7 +225,7 @@ class ElasticKafkaApis( case None => (-1, -1) } } - LeaderNode(leaderId, leaderEpoch, OptionConverters.toScala(produceRouter.getLeaderNode(leaderId, clientIdMetadata, ln.value()))) + LeaderNode(leaderId, leaderEpoch, OptionConverters.toScala(trafficInterceptor.getLeaderNode(leaderId, clientIdMetadata, ln.value()))) } /** @@ -259,7 +274,7 @@ class ElasticKafkaApis( } }) - val clientIdMetadata = ClientIdMetadata.of(request.header.clientId(), request.context.clientAddress) + val clientIdMetadata = ClientIdMetadata.of(request.header.clientId(), request.context.clientAddress, request.context.connectionId) // the callback for sending a produce response // The construction of ProduceResponse is able to accept auto-generated protocol data so @@ -286,7 +301,7 @@ class ElasticKafkaApis( if (request.header.apiVersion >= 10) { status.error match { case Errors.NOT_LEADER_OR_FOLLOWER => - val leaderNode = getCurrentLeaderForProduce(topicPartition, clientIdMetadata, request.context.listenerName) + val leaderNode = getCurrentLeaderForProduceAndFetch(topicPartition, request.context.listenerName, clientIdMetadata) leaderNode.node.foreach { node => nodeEndpoints.put(node.id(), node) } @@ -311,9 +326,7 @@ class ElasticKafkaApis( val requestThrottleTimeMs = if (produceRequest.acks == 0) 0 else quotas.request.maybeRecordAndGetThrottleTimeMs(request, timeMs) - val brokerRequestThrottleTimeMs = - if (produceRequest.acks == 0) 0 - else quotas.broker.maybeRecordAndGetThrottleTimeMs(QuotaType.Request, request, 1, timeMs) + val brokerRequestThrottleTimeMs = quotas.broker.maybeRecordAndGetThrottleTimeMs(QuotaType.RequestRate, request, 1, timeMs) val maxThrottleTimeMs = IntStream.of(bandwidthThrottleTimeMs, requestThrottleTimeMs, brokerBandwidthThrottleTimeMs, brokerRequestThrottleTimeMs).max().orElse(0) if (maxThrottleTimeMs > 0) { request.apiThrottleTimeMs = maxThrottleTimeMs @@ -324,7 +337,7 @@ class ElasticKafkaApis( } else if (brokerBandwidthThrottleTimeMs == maxThrottleTimeMs) { requestHelper.throttle(QuotaType.Produce, quotas.broker, request, brokerBandwidthThrottleTimeMs) } else if (brokerRequestThrottleTimeMs == maxThrottleTimeMs) { - requestHelper.throttle(QuotaType.Request, quotas.broker, request, brokerRequestThrottleTimeMs) + requestHelper.throttle(QuotaType.RequestRate, quotas.broker, request, brokerRequestThrottleTimeMs) } } // AutoMQ for Kafka inject end @@ -386,16 +399,20 @@ class ElasticKafkaApis( } def doAppendRecords(): Unit = { - produceRouter.handleProduceRequest( - request.header.apiVersion, - clientIdMetadata, - produceRequest.timeout, - produceRequest.acks, - internalTopicsAllowed, - produceRequest.transactionalId, - authorizedRequestInfo.asJava, - sendResponseCallbackJava, - processingStatsCallbackJava, + GlobalNetworkBandwidthLimiters.instance().inbound().consume(ThrottleStrategy.BYPASS, request.sizeInBytes) + trafficInterceptor.handleProduceRequest( + ProduceRequestArgs.builder() + .apiVersion(request.header.apiVersion) + .clientId(clientIdMetadata) + .timeout(produceRequest.timeout) + .requiredAcks(produceRequest.acks) + .internalTopicsAllowed(internalTopicsAllowed) + .transactionId(produceRequest.transactionalId) + .entriesPerPartition(authorizedRequestInfo.asJava) + .responseCallback(sendResponseCallbackJava) + .recordValidationStatsCallback(processingStatsCallbackJava) + .requestLocal(requestLocal) + .build() ) // if the request is put into the purgatory, it will have a held reference and hence cannot be garbage collected; @@ -409,9 +426,21 @@ class ElasticKafkaApis( } } + def handleUpdateGroupRequest(request: RequestChannel.Request, requestLocal: RequestLocal): Unit = { + val updateGroupsRequest = request.body[AutomqUpdateGroupRequest] + groupCoordinator.updateGroup(request.context, updateGroupsRequest.data(), requestLocal.bufferSupplier) + .whenComplete((response, ex) => { + if (ex != null) { + requestHelper.sendMaybeThrottle(request, updateGroupsRequest.getErrorResponse(ex)) + } else { + requestHelper.sendMaybeThrottle(request, new AutomqUpdateGroupResponse(response)) + } + }) + } + def handleZoneRouterRequest(request: RequestChannel.Request, requestLocal: RequestLocal): Unit = { val zoneRouterRequest = request.body[AutomqZoneRouterRequest] - produceRouter.handleZoneRouterRequest(zoneRouterRequest.data().metadata()).thenAccept(response => { + trafficInterceptor.handleZoneRouterRequest(zoneRouterRequest.data()).thenAccept(response => { requestChannel.sendResponse(request, response, None) }).exceptionally(ex => { handleError(request, ex) @@ -419,24 +448,20 @@ class ElasticKafkaApis( }) } - def handleProduceAppendJavaCompatible(timeout: Long, - requiredAcks: Short, - internalTopicsAllowed: Boolean, - transactionalId: String, - entriesPerPartition: util.Map[TopicPartition, MemoryRecords], - responseCallback: util.Map[TopicPartition, PartitionResponse] => Unit, - recordValidationStatsCallback: util.Map[TopicPartition, RecordValidationStats] => Unit = _ => (), - apiVersion: Short): Unit = { - val transactionSupportedOperation = if (apiVersion > 10) genericError else defaultError + def handleProduceAppendJavaCompatible( + args: ProduceRequestArgs, + ): Unit = { + val transactionSupportedOperation = if (args.apiVersion() > 10) genericError else defaultError replicaManager.handleProduceAppend( - timeout = timeout, - requiredAcks = requiredAcks, - internalTopicsAllowed = internalTopicsAllowed, - transactionalId = transactionalId, - entriesPerPartition = entriesPerPartition.asScala, - responseCallback = rst => responseCallback.apply(rst.asJava), - recordValidationStatsCallback = rst => recordValidationStatsCallback.apply(rst.asJava), - transactionSupportedOperation = transactionSupportedOperation + timeout = args.timeout(), + requiredAcks = args.requiredAcks(), + internalTopicsAllowed = args.internalTopicsAllowed(), + transactionalId = args.transactionId(), + entriesPerPartition = args.entriesPerPartition().asScala, + responseCallback = rst => args.responseCallback().accept(rst.asJava), + recordValidationStatsCallback = rst => args.recordValidationStatsCallback().accept(rst.asJava), + transactionSupportedOperation = transactionSupportedOperation, + requestLocal = args.requestLocal(), ) } @@ -591,6 +616,7 @@ class ElasticKafkaApis( val partitions = new util.LinkedHashMap[TopicIdPartition, FetchResponseData.PartitionData] val reassigningPartitions = mutable.Set[TopicIdPartition]() val nodeEndpoints = new mutable.HashMap[Int, Node] + val clientIdMetadata = ClientIdMetadata.of(request.header.clientId(), request.context.clientAddress, request.context.connectionId) responsePartitionData.foreach { case (tp, data) => val abortedTransactions = data.abortedTransactions.orElse(null) val lastStableOffset: Long = data.lastStableOffset.orElse(FetchResponse.INVALID_LAST_STABLE_OFFSET) @@ -608,7 +634,7 @@ class ElasticKafkaApis( if (versionId >= 16) { data.error match { case Errors.NOT_LEADER_OR_FOLLOWER | Errors.FENCED_LEADER_EPOCH => - val leaderNode = getCurrentLeader(tp.topicPartition(), request.context.listenerName) + val leaderNode = getCurrentLeaderForProduceAndFetch(tp.topicPartition(), request.context.listenerName, clientIdMetadata) leaderNode.node.foreach { node => nodeEndpoints.put(node.id(), node) } @@ -697,26 +723,34 @@ class ElasticKafkaApis( val timeMs = time.milliseconds() // AutoMQ for Kafka inject start + val isSlowRead = !ReadHint.isFastRead + val requestThrottleTimeMs = quotas.request.maybeRecordAndGetThrottleTimeMs(request, timeMs) val bandwidthThrottleTimeMs = quotas.fetch.maybeRecordAndGetThrottleTimeMs(request, responseSize, timeMs) - val brokerRequestThrottleTimeMs = quotas.broker.maybeRecordAndGetThrottleTimeMs(QuotaType.Request, request, 1, timeMs) val brokerBandwidthThrottleTimeMs = quotas.broker.maybeRecordAndGetThrottleTimeMs(QuotaType.Fetch, request, responseSize, timeMs) + val brokerSlowFetchThrottleTimeMs = if (isSlowRead) quotas.broker.maybeRecordAndGetThrottleTimeMs(QuotaType.SlowFetch, request, responseSize, timeMs) else 0 + val brokerRequestThrottleTimeMs = quotas.broker.maybeRecordAndGetThrottleTimeMs(QuotaType.RequestRate, request, 1, timeMs) - val maxThrottleTimeMs = IntStream.of(bandwidthThrottleTimeMs, requestThrottleTimeMs, brokerBandwidthThrottleTimeMs, brokerRequestThrottleTimeMs).max().orElse(0) + val maxThrottleTimeMs = IntStream.of(bandwidthThrottleTimeMs, requestThrottleTimeMs, brokerBandwidthThrottleTimeMs, brokerSlowFetchThrottleTimeMs, brokerRequestThrottleTimeMs).max().orElse(0) if (maxThrottleTimeMs > 0) { request.apiThrottleTimeMs = maxThrottleTimeMs // Even if we need to throttle for request quota violation, we should "unrecord" the already recorded value // from the fetch quota because we are going to return an empty response. quotas.fetch.unrecordQuotaSensor(request, responseSize, timeMs) quotas.broker.unrecordQuotaSensor(QuotaType.Fetch, responseSize, timeMs) + if (isSlowRead) { + quotas.broker.unrecordQuotaSensor(QuotaType.SlowFetch, responseSize, timeMs) + } if (bandwidthThrottleTimeMs == maxThrottleTimeMs) { requestHelper.throttle(quotas.fetch, request, bandwidthThrottleTimeMs) } else if (requestThrottleTimeMs == maxThrottleTimeMs) { requestHelper.throttle(quotas.request, request, requestThrottleTimeMs) } else if (brokerBandwidthThrottleTimeMs == maxThrottleTimeMs) { requestHelper.throttle(QuotaType.Fetch, quotas.broker, request, brokerBandwidthThrottleTimeMs) + } else if (brokerSlowFetchThrottleTimeMs == maxThrottleTimeMs) { + requestHelper.throttle(QuotaType.SlowFetch, quotas.broker, request, brokerSlowFetchThrottleTimeMs) } else if (brokerRequestThrottleTimeMs == maxThrottleTimeMs) { - requestHelper.throttle(QuotaType.Request, quotas.broker, request, brokerRequestThrottleTimeMs) + requestHelper.throttle(QuotaType.RequestRate, quotas.broker, request, brokerRequestThrottleTimeMs) } // AutoMQ for Kafka inject end @@ -746,7 +780,7 @@ class ElasticKafkaApis( Int.MaxValue else { val maxValue = quotas.fetch.getMaxValueInQuotaWindow(request.session, clientId).toInt - val brokerMaxValue = quotas.broker.getMaxValueInQuotaWindow(QuotaType.Fetch).toInt + val brokerMaxValue = quotas.broker.getMaxValueInQuotaWindow(QuotaType.Fetch, request).toInt math.min(maxValue, brokerMaxValue) } @@ -792,12 +826,35 @@ class ElasticKafkaApis( listOffsetHandleExecutor.execute(() => super.handleListOffsetRequest(request)) } - override protected def metadataTopicsInterceptor(clientId: String, listenerName: String, topics: util.List[MetadataResponseData.MetadataResponseTopic]): util.List[MetadataResponseData.MetadataResponseTopic] = { - produceRouter.handleMetadataResponse(clientId, topics) + override def handleOffsetForLeaderEpochRequest(request: RequestChannel.Request): Unit = { + val cf = snapshotAwaitReadySupplier.get() + offsetForLeaderEpochExecutor.execute(() => { + // Await new snapshots to be applied to avoid consumers finding the endOffset jumping back when the snapshot-read partition leader changes. + cf.join() + super.handleOffsetForLeaderEpochRequest(request) + }) + } + + override protected def metadataTopicsInterceptor(clientId: ClientIdMetadata, listenerName: String, topics: util.List[MetadataResponseData.MetadataResponseTopic]): util.List[MetadataResponseData.MetadataResponseTopic] = { + trafficInterceptor.handleMetadataResponse(clientId, topics) + } + + def handleGetPartitionSnapshotRequest(request: RequestChannel.Request, requestLocal: RequestLocal): Unit = { + val req = request.body[AutomqGetPartitionSnapshotRequest] + replicaManager.asInstanceOf[ElasticReplicaManager].handleGetPartitionSnapshotRequest(req) + .thenAccept(resp => requestHelper.sendMaybeThrottle(request, resp)) + .exceptionally(ex => { + handleError(request, ex) + null + }) + } + + def setTrafficInterceptor(trafficInterceptor: TrafficInterceptor): Unit = { + this.trafficInterceptor = trafficInterceptor } - def setProduceRouter(produceRouter: ProduceRouter): Unit = { - this.produceRouter = produceRouter + def setSnapshotAwaitReadyProvider(supplier: Supplier[CompletableFuture[Void]]): Unit = { + this.snapshotAwaitReadySupplier = supplier } } diff --git a/core/src/main/scala/kafka/server/streamaspect/ElasticReplicaManager.scala b/core/src/main/scala/kafka/server/streamaspect/ElasticReplicaManager.scala index 966efbcfba..250a9808dd 100644 --- a/core/src/main/scala/kafka/server/streamaspect/ElasticReplicaManager.scala +++ b/core/src/main/scala/kafka/server/streamaspect/ElasticReplicaManager.scala @@ -1,10 +1,17 @@ package kafka.server.streamaspect import com.automq.stream.api.exceptions.FastReadFailFastException -import com.automq.stream.utils.FutureUtil +import com.automq.stream.s3.metrics.{MetricsLevel, TimerUtil} +import com.automq.stream.s3.network.{AsyncNetworkBandwidthLimiter, GlobalNetworkBandwidthLimiters, ThrottleStrategy} +import com.automq.stream.utils.{FutureUtil, Systems} +import com.automq.stream.utils.threads.S3StreamThreadPoolMonitor +import kafka.automq.interceptor.{ClientIdKey, ClientIdMetadata, TrafficInterceptor} +import kafka.automq.kafkalinking.KafkaLinkingManager +import kafka.automq.partition.snapshot.PartitionSnapshotsManager +import kafka.automq.zerozone.ZeroZoneThreadLocalContext import kafka.cluster.Partition import kafka.log.remote.RemoteLogManager -import kafka.log.streamaspect.{ElasticLogManager, PartitionStatusTracker, ReadHint} +import kafka.log.streamaspect.{ElasticLogManager, OpenHint, PartitionStatusTracker, ReadHint} import kafka.log.{LogManager, UnifiedLog} import kafka.server.Limiter.Handler import kafka.server.QuotaFactory.QuotaManagers @@ -14,14 +21,17 @@ import kafka.server.checkpoints.{LazyOffsetCheckpoints, OffsetCheckpoints} import kafka.utils.Implicits.MapExtensionMethods import kafka.utils.{CoreUtils, Exit} import kafka.zk.KafkaZkClient +import org.apache.commons.lang3.StringUtils import org.apache.kafka.common.errors._ import org.apache.kafka.common.errors.s3.StreamFencedException import org.apache.kafka.common.internals.Topic import org.apache.kafka.common.metrics.Metrics import org.apache.kafka.common.protocol.Errors import org.apache.kafka.common.record.{MemoryRecords, PooledRecords, PooledResource} +import org.apache.kafka.common.replica.ClientMetadata import org.apache.kafka.common.requests.FetchRequest import org.apache.kafka.common.requests.FetchRequest.PartitionData +import org.apache.kafka.common.requests.s3.{AutomqGetPartitionSnapshotRequest, AutomqGetPartitionSnapshotResponse} import org.apache.kafka.common.utils.{ThreadUtils, Time} import org.apache.kafka.common.{TopicIdPartition, TopicPartition, Uuid} import org.apache.kafka.image.{LocalReplicaChanges, MetadataImage, TopicsDelta} @@ -36,11 +46,12 @@ import java.util import java.util.Optional import java.util.concurrent._ import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong, AtomicReference} -import java.util.function.Consumer +import java.util.function.{BiFunction, Consumer} import scala.collection.mutable.ArrayBuffer import scala.collection.{Seq, mutable} import scala.compat.java8.OptionConverters -import scala.jdk.CollectionConverters.{CollectionHasAsScala, MapHasAsScala, SetHasAsJava} +import scala.compat.java8.OptionConverters.RichOptionalGeneric +import scala.jdk.CollectionConverters.{CollectionHasAsScala, EnumerationHasAsScala, MapHasAsScala, SetHasAsJava} object ElasticReplicaManager { def emptyReadResults(partitions: Seq[TopicIdPartition]): Seq[(TopicIdPartition, LogReadResult)] = { @@ -83,8 +94,8 @@ class ElasticReplicaManager( brokerEpochSupplier: () => Long = () => -1, addPartitionsToTxnManager: Option[AddPartitionsToTxnManager] = None, directoryEventHandler: DirectoryEventHandler = DirectoryEventHandler.NOOP, - private val fastFetchExecutor: ExecutorService = Executors.newFixedThreadPool(4, ThreadUtils.createThreadFactory("kafka-apis-fast-fetch-executor-%d", true)), - private val slowFetchExecutor: ExecutorService = Executors.newFixedThreadPool(12, ThreadUtils.createThreadFactory("kafka-apis-slow-fetch-executor-%d", true)), + private val fastFetchExecutor: ExecutorService = S3StreamThreadPoolMonitor.createAndMonitor(4, 4, 0L, TimeUnit.MILLISECONDS, "kafka-apis-fast-fetch-executor", true, 10000), + private val slowFetchExecutor: ExecutorService = S3StreamThreadPoolMonitor.createAndMonitor(12, 12, 0L, TimeUnit.MILLISECONDS, "kafka-apis-slow-fetch-executor", true, 10000), private val partitionMetricsCleanerExecutor: ScheduledExecutorService = Executors.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("kafka-partition-metrics-cleaner", true)), ) extends ReplicaManager(config, metrics, time, scheduler, logManager, remoteLogManager, quotaManagers, metadataCache, logDirFailureChannel, alterPartitionManager, brokerTopicStats, isShuttingDown, zkClient, delayedProducePurgatoryParam, @@ -93,7 +104,7 @@ class ElasticReplicaManager( directoryEventHandler) { partitionMetricsCleanerExecutor.scheduleAtFixedRate(() => { - brokerTopicStats.removeRedundantMetrics(allPartitions.keys) + brokerTopicStats.removeRedundantMetrics(allPartitions.keys ++ snapshotReadPartitions.keys.asScala) }, 1, 1, TimeUnit.HOURS) protected val openingPartitions = new ConcurrentHashMap[TopicPartition, CompletableFuture[Void]]() @@ -113,14 +124,32 @@ class ElasticReplicaManager( fetchExecutorQueueSizeGaugeMap }) - private val fastFetchLimiter = new FairLimiter(200 * 1024 * 1024) // 200MiB - private val slowFetchLimiter = new FairLimiter(200 * 1024 * 1024) // 200MiB - private val fetchLimiterGaugeMap = new util.HashMap[String, Integer]() + private val fetchLimiterSize = Systems.getEnvInt("AUTOMQ_FETCH_LIMITER_SIZE", + // autoscale the fetch limiter size based on heap size, min 200MiB, max 1GiB, every 3GB heap add 100MiB limiter + Math.min(1024, 100 * Math.max(2, (Systems.HEAP_MEMORY_SIZE / (1024 * 1024 * 1024) / 3)).asInstanceOf[Int]) * 1024 * 1024 + ) + private val fastFetchLimiter = new FairLimiter(fetchLimiterSize, FETCH_LIMITER_FAST_NAME) + private val slowFetchLimiter = new FairLimiter(fetchLimiterSize, FETCH_LIMITER_SLOW_NAME) + private val fetchLimiterWaitingTasksGaugeMap = new util.HashMap[String, Integer]() + S3StreamKafkaMetricsManager.setFetchLimiterWaitingTaskNumSupplier(() => { + fetchLimiterWaitingTasksGaugeMap.put(FETCH_LIMITER_FAST_NAME, fastFetchLimiter.waitingThreads()) + fetchLimiterWaitingTasksGaugeMap.put(FETCH_LIMITER_SLOW_NAME, slowFetchLimiter.waitingThreads()) + fetchLimiterWaitingTasksGaugeMap + }) + private val fetchLimiterPermitsGaugeMap = new util.HashMap[String, Integer]() S3StreamKafkaMetricsManager.setFetchLimiterPermitNumSupplier(() => { - fetchLimiterGaugeMap.put(FETCH_LIMITER_FAST_NAME, fastFetchLimiter.availablePermits()) - fetchLimiterGaugeMap.put(FETCH_LIMITER_SLOW_NAME, slowFetchLimiter.availablePermits()) - fetchLimiterGaugeMap + fetchLimiterPermitsGaugeMap.put(FETCH_LIMITER_FAST_NAME, fastFetchLimiter.availablePermits()) + fetchLimiterPermitsGaugeMap.put(FETCH_LIMITER_SLOW_NAME, slowFetchLimiter.availablePermits()) + fetchLimiterPermitsGaugeMap }) + private val fetchLimiterTimeoutCounterMap = util.Map.of( + fastFetchLimiter.name, S3StreamKafkaMetricsManager.buildFetchLimiterTimeoutMetric(fastFetchLimiter.name), + slowFetchLimiter.name, S3StreamKafkaMetricsManager.buildFetchLimiterTimeoutMetric(slowFetchLimiter.name) + ) + private val fetchLimiterTimeHistogramMap = util.Map.of( + fastFetchLimiter.name, S3StreamKafkaMetricsManager.buildFetchLimiterTimeMetric(MetricsLevel.INFO, fastFetchLimiter.name), + slowFetchLimiter.name, S3StreamKafkaMetricsManager.buildFetchLimiterTimeMetric(MetricsLevel.INFO, slowFetchLimiter.name) + ) /** * Used to reduce allocation in [[readFromLocalLogV2]] @@ -164,6 +193,22 @@ class ElasticReplicaManager( private var fenced: Boolean = false + private val partitionLifecycleListeners = new util.ArrayList[PartitionLifecycleListener]() + + private var kafkaLinkingManager = Option.empty[KafkaLinkingManager] + + private var partitionSnapshotsManager: PartitionSnapshotsManager = null + + private val snapshotReadPartitions = new ConcurrentHashMap[TopicPartition, Partition]() + + private var trafficInterceptor: TrafficInterceptor = null + + addPartitionLifecycleListener(new PartitionLifecycleListener { + override def onOpen(partition: Partition): Unit = partitionSnapshotsManager.onPartitionOpen(partition) + + override def onClose(partition: Partition): Unit = partitionSnapshotsManager.onPartitionClose(partition) + }) + override def startup(): Unit = { super.startup() val haltBrokerOnFailure = metadataCache.metadataVersion().isLessThan(MetadataVersion.IBP_1_0_IV0) @@ -185,6 +230,7 @@ class ElasticReplicaManager( val partitions = partitionsToStop.map(_.topicPartition) replicaFetcherManager.removeFetcherForPartitions(partitions) replicaAlterLogDirsManager.removeFetcherForPartitions(partitions) + kafkaLinkingManager.foreach(_.removePartitions(partitions.asJava)) // Second remove deleted partitions from the partition map. Fetchers rely on the // ReplicaManager to get Partition's information so they must be stopped first. @@ -195,6 +241,7 @@ class ElasticReplicaManager( getPartition(topicPartition) match { case hostedPartition: HostedPartition.Online => if (allPartitions.remove(topicPartition, hostedPartition)) { + notifyPartitionClose(hostedPartition.partition) brokerTopicStats.removeMetrics(topicPartition) maybeRemoveTopicMetrics(topicPartition.topic) // AutoMQ for Kafka inject start @@ -205,10 +252,6 @@ class ElasticReplicaManager( val start = System.currentTimeMillis() hostedPartition.partition.close() info(s"partition $topicPartition is closed, cost ${System.currentTimeMillis() - start} ms") - if (!metadataCache.autoMQVersion().isReassignmentV1Supported) { - // TODO: https://github.com/AutoMQ/automq/issues/1153 add schedule check when leader isn't successfully set - alterPartitionManager.tryElectLeader(topicPartition) - } } else { // Logs are not deleted here. They are deleted in a single batch later on. // This is done to avoid having to checkpoint for every deletions. @@ -252,8 +295,14 @@ class ElasticReplicaManager( /** * Remove the usage of [[Option]] in [[getPartition]] to avoid allocation */ - def getPartitionV2(topicPartition: TopicPartition): HostedPartition = { - val partition = allPartitions.get(topicPartition) + override def getPartition(topicPartition: TopicPartition): HostedPartition = { + var partition = allPartitions.get(topicPartition) + if (partition == null) { + val p = snapshotReadPartitions.get(topicPartition) + if (p != null) { + partition = HostedPartition.Online(p) + } + } if (null == partition) { HostedPartition.None } else { @@ -261,11 +310,20 @@ class ElasticReplicaManager( } } + def getPartitionWithoutSnapshotRead(topicPartition: TopicPartition): HostedPartition = { + val partition = allPartitions.get(topicPartition) + if (partition == null) { + HostedPartition.None + } else { + partition + } + } + /** * Remove the usage of [[Either]] in [[getPartitionOrException]] to avoid allocation */ def getPartitionOrExceptionV2(topicPartition: TopicPartition): Partition = { - getPartitionV2(topicPartition) match { + getPartition(topicPartition) match { case HostedPartition.Online(partition) => partition case HostedPartition.Offline(partition) => @@ -311,7 +369,6 @@ class ElasticReplicaManager( entriesPerPartition.map { case (topicPartition, records) => brokerTopicStats.topicStats(topicPartition.topic).totalProduceRequestRate.mark() - brokerTopicStats.topicPartitionStats(topicPartition).totalProduceRequestRate.mark() brokerTopicStats.allTopicsStats.totalProduceRequestRate.mark() // reject appending to internal topics if it is not allowed @@ -328,7 +385,7 @@ class ElasticReplicaManager( val numAppendedMessages = info.numMessages // update stats for successfully appended bytes and messages as bytesInRate and messageInRate - brokerTopicStats.topicPartitionStats(topicPartition).bytesInRate.mark(records.sizeInBytes()) + brokerTopicStats.updatePartitionBytesIn(topicPartition, records.sizeInBytes()) brokerTopicStats.topicStats(topicPartition.topic).bytesInRate.mark(records.sizeInBytes) brokerTopicStats.allTopicsStats.bytesInRate.mark(records.sizeInBytes) brokerTopicStats.topicStats(topicPartition.topic).messagesInRate.mark(numAppendedMessages) @@ -450,7 +507,7 @@ class ElasticReplicaManager( logReadResults.foreach { case (topicIdPartition, logReadResult) => brokerTopicStats.topicStats(topicIdPartition.topicPartition.topic).totalFetchRequestRate.mark() - brokerTopicStats.topicPartitionStats(topicIdPartition.topicPartition).totalFetchRequestRate.mark() + brokerTopicStats.updatePartitionFetchRequestRate(topicIdPartition.topicPartition()) brokerTopicStats.allTopicsStats.totalFetchRequestRate.mark() if (logReadResult.error != Errors.NONE) errorReadingData = true @@ -557,15 +614,17 @@ class ElasticReplicaManager( math.min(bytesNeedFromParam, limiter.maxPermits()) } + val timer: TimerUtil = new TimerUtil() val handler: Handler = timeoutMs match { case t if t > 0 => limiter.acquire(bytesNeed(), t) case _ => limiter.acquire(bytesNeed()) } + fetchLimiterTimeHistogramMap.get(limiter.name).record(timer.elapsedAs(TimeUnit.NANOSECONDS)) if (handler == null) { - // handler maybe null if it timed out to acquire from limiter - // TODO add metrics for this - // warn(s"Returning emtpy fetch response for fetch request $readPartitionInfo since the wait time exceeds $timeoutMs ms.") + // the handler will be null if it timed out to acquire from limiter + fetchLimiterTimeoutCounterMap.get(limiter.name).add(MetricsLevel.INFO, 1) + // warn(s"Returning empty fetch response for fetch request $readPartitionInfo since the wait time exceeds $timeoutMs ms.") ElasticReplicaManager.emptyReadResults(readPartitionInfo.map(_._1)) } else { try { @@ -600,6 +659,34 @@ class ElasticReplicaManager( } } + def getSnapshotReadPreferredNode(clientMetadataOpt: Optional[ClientMetadata]): Optional[Int] = { + if (clientMetadataOpt.isPresent && trafficInterceptor != null) { + val clientMetadata = clientMetadataOpt.get() + val clientIdMetadata = ClientIdMetadata.of(clientMetadata.clientId(), clientMetadata.clientAddress(), null) + if (StringUtils.isNotBlank(clientMetadata.rackId())) { + clientIdMetadata.metadata(ClientIdKey.AVAILABILITY_ZONE, util.List.of(clientMetadata.rackId())) + } + val nodeOpt = trafficInterceptor.getLeaderNode(config.nodeId, clientIdMetadata, clientMetadata.listenerName()) + if (StringUtils.isBlank(clientMetadata.rackId())) { + if (nodeOpt.isPresent && nodeOpt.get().id() != config.nodeId) { + // the consumer should directly read from the snapshot-read partition + Optional.of(-1) + } else { + Optional.empty() + } + } else { + if (nodeOpt.isPresent && nodeOpt.get().id() != config.nodeId) { + // return the preferred node + Optional.of(nodeOpt.get().id()) + } else { + Optional.empty() + } + } + } else { + Optional.empty() + } + } + /** * Parallel read from multiple topic partitions at the given offset up to maxSize bytes */ @@ -612,6 +699,9 @@ class ElasticReplicaManager( val fastReadFastFail = new AtomicReference[FastReadFailFastException]() + // snapshot-read preferred node + val snapshotReadPreferredNode = getSnapshotReadPreferredNode(params.clientMetadata) + /** * Convert a throwable to [[LogReadResult]] with [[LogReadResult.exception]] set. * Note: All parameters except `throwable` are just used for logging or metrics. @@ -727,10 +817,22 @@ class ElasticReplicaManager( val fetchTimeMs = time.milliseconds - // ~~ If we are the leader, determine the preferred read-replica ~~ - // NOTE: We do not check the preferred read-replica like Apache Kafka does in - // [[ReplicaManager.readFromLocalLog]], as we always have only one replica per partition. - val preferredReadReplica = None + val preferredReadReplica = snapshotReadPreferredNode.asScala + if (preferredReadReplica.isDefined) { + // If a preferred read-replica is set, skip the read + val offsetSnapshot = partition.fetchOffsetSnapshot(fetchInfo.currentLeaderEpoch, fetchOnlyFromLeader = false) + val rst = LogReadResult(info = new FetchDataInfo(LogOffsetMetadata.UNKNOWN_OFFSET_METADATA, MemoryRecords.EMPTY), + divergingEpoch = None, + highWatermark = offsetSnapshot.highWatermark.messageOffset, + leaderLogStartOffset = offsetSnapshot.logStartOffset, + leaderLogEndOffset = offsetSnapshot.logEndOffset.messageOffset, + followerLogStartOffset = fetchInfo.logStartOffset, + fetchTimeMs = -1L, + lastStableOffset = Some(offsetSnapshot.lastStableOffset.messageOffset), + preferredReadReplica = preferredReadReplica, + exception = None) + return CompletableFuture.completedFuture(rst) + } // Try the read first, this tells us whether we need all of adjustedFetchSize for this partition partition.fetchRecordsAsync( @@ -788,19 +890,24 @@ class ElasticReplicaManager( } var partitionIndex = 0; - while (remainingBytes.get() > 0 && partitionIndex < readPartitionInfo.size) { + while (remainingBytes.get() > 0 && partitionIndex < readPartitionInfo.size && fastReadFastFail.get() == null) { // In each iteration, we read as many partitions as possible until we reach the maximum bytes limit. val readCfArray = readFutureBuffer.get() readCfArray.clear() var assignedBytes = 0 // The total bytes we have assigned to the read requests. val availableBytes = remainingBytes.get() // The remaining bytes we can assign to the read requests, used to control the following loop. - while (assignedBytes < availableBytes && partitionIndex < readPartitionInfo.size) { + while (assignedBytes < availableBytes && partitionIndex < readPartitionInfo.size + // When there is a fast read exception, quit the loop earlier. + && fastReadFastFail.get() == null) { // Iterate over the partitions. val tp = readPartitionInfo(partitionIndex)._1 val partitionData = readPartitionInfo(partitionIndex)._2 try { val partition = getPartitionAndCheckTopicId(tp) + if (snapshotReadPreferredNode.isPresent && snapshotReadPreferredNode.get() == -1) { + throw new NotLeaderOrFollowerException("The consumer should read the snapshot-read partition in the same rack") + } val logReadInfo = partition.checkFetchOffsetAndMaybeGetInfo(params, partitionData) if (null != logReadInfo) { @@ -867,9 +974,15 @@ class ElasticReplicaManager( release() throw fastReadFastFail.get() } + acquireNetworkOutPermit(limitBytes - remainingBytes.get(), if (ReadHint.isFastRead) ThrottleStrategy.TAIL else ThrottleStrategy.CATCH_UP) result } + private def acquireNetworkOutPermit(size: Int, throttleStrategy: ThrottleStrategy): Unit = { + GlobalNetworkBandwidthLimiters.instance().get(AsyncNetworkBandwidthLimiter.Type.OUTBOUND) + .consume(throttleStrategy, size).join() + } + def handlePartitionFailure(partitionDir: String): Unit = { warn(s"Stopping serving partition $partitionDir") replicaStateChangeLock synchronized { @@ -887,6 +1000,7 @@ class ElasticReplicaManager( replicaFetcherManager.removeFetcherForPartitions(newOfflinePartitions) replicaAlterLogDirsManager.removeFetcherForPartitions(newOfflinePartitions ++ partitionsWithOfflineFutureReplica.map(_.topicPartition)) + kafkaLinkingManager.foreach(_.removePartitions(newOfflinePartitions.asJava)) // These partitions should first be made offline to remove topic metrics. newOfflinePartitions.foreach { topicPartition => @@ -912,7 +1026,7 @@ class ElasticReplicaManager( delta: TopicsDelta, topicId: Uuid, createHook: Consumer[Partition] = _ => {}): Option[(Partition, Boolean)] = { - getPartition(tp) match { + getPartitionWithoutSnapshotRead(tp) match { case HostedPartition.Offline(offlinePartition) => if (offlinePartition.flatMap(p => p.topicId).contains(topicId)) { stateChangeLogger.warn(s"Unable to bring up new local leader $tp " + @@ -955,6 +1069,7 @@ class ElasticReplicaManager( val partition = Partition(tp, time, this) createHook.accept(partition) allPartitions.put(tp, HostedPartition.Online(partition)) + notifyPartitionOpen(partition) Some(partition, true) } } @@ -1186,6 +1301,7 @@ class ElasticReplicaManager( }).foreach { case (partition, _) => try { changedPartitions.add(partition) + kafkaLinkingManager.foreach(_.addPartitions(Set(partition.topicPartition).asJava)) } catch { case e: KafkaStorageException => stateChangeLogger.info(s"Skipped the become-leader state change for $tp " + @@ -1276,7 +1392,6 @@ class ElasticReplicaManager( } partitionOpenOpExecutor.shutdown() partitionCloseOpExecutor.shutdown() - CoreUtils.swallow(ElasticLogManager.shutdown(), this) } /** @@ -1320,7 +1435,7 @@ class ElasticReplicaManager( transactionWaitingForValidationMap.computeIfAbsent(producerId, _ => { Verification( new AtomicBoolean(false), - new ArrayBlockingQueue[TransactionVerificationRequest](5), + new LinkedBlockingQueue[TransactionVerificationRequest](), new AtomicLong(time.milliseconds())) }) } else { @@ -1348,25 +1463,31 @@ class ElasticReplicaManager( verification: Verification, callback: (RequestLocal, T) => Unit, ): (RequestLocal, T) => Unit = { + val writeContext = ZeroZoneThreadLocalContext.writeContext().detach() (requestLocal, args) => { try { + // The thread switch, so we need to attach the write context to the current thread. + ZeroZoneThreadLocalContext.attach(writeContext) callback(requestLocal, args) } catch { case e: Throwable => error("Error in transaction verification callback", e) } if (verification != null) { + var request: TransactionVerificationRequest = null verification.synchronized { verification.timestamp.set(time.milliseconds()) if (!verification.waitingRequests.isEmpty) { // Since the callback thread and task thread may be different, we need to ensure that the tasks are executed sequentially. - val request = verification.waitingRequests.poll() - request.task() + request = verification.waitingRequests.poll() } else { // If there are no tasks in the queue, set hasInflight to false verification.hasInflight.set(false) } } + if (request != null) { + request.task() + } val lastCleanTimestamp = lastTransactionCleanTimestamp.get(); val now = time.milliseconds() if (now - lastCleanTimestamp > 60 * 1000 && lastTransactionCleanTimestamp.compareAndSet(lastCleanTimestamp, now)) { @@ -1382,4 +1503,62 @@ class ElasticReplicaManager( } } } + + def handleGetPartitionSnapshotRequest(request: AutomqGetPartitionSnapshotRequest): CompletableFuture[AutomqGetPartitionSnapshotResponse] = { + partitionSnapshotsManager.handle(request) + } + + def addPartitionLifecycleListener(listener: PartitionLifecycleListener): Unit = { + partitionLifecycleListeners.add(listener) + } + + def computeSnapshotReadPartition(topicPartition: TopicPartition, + remappingFunction: BiFunction[TopicPartition, Partition, Partition]): Partition = { + snapshotReadPartitions.compute(topicPartition, (tp, partition) => { + val newPartition = remappingFunction.apply(tp, partition) + if (newPartition == null) { + brokerTopicStats.removeMetrics(tp) + } + newPartition + }) + } + + def newSnapshotReadPartition(topicIdPartition: TopicIdPartition): Partition = { + OpenHint.markSnapshotRead() + val partition = Partition.apply(topicIdPartition, time, this) + partition.leaderReplicaIdOpt = Some(localBrokerId) + partition.createLogIfNotExists(true, false, new LazyOffsetCheckpoints(highWatermarkCheckpoints), partition.topicId, Option.empty) + OpenHint.clear() + partition + } + + private def notifyPartitionOpen(partition: Partition): Unit = { + partitionLifecycleListeners.forEach(listener => CoreUtils.swallow(listener.onOpen(partition), this)) + } + + private def notifyPartitionClose(partition: Partition): Unit = { + partitionLifecycleListeners.forEach(listener => CoreUtils.swallow(listener.onClose(partition), this)) + } + + override def shutdown(checkpointHW: Boolean): Unit = { + kafkaLinkingManager.foreach(_.shutdown()) + super.shutdown(checkpointHW) + } + + def setKafkaLinkingManager(kafkaLinkingManager: KafkaLinkingManager): Unit = { + if (kafkaLinkingManager == null) { + this.kafkaLinkingManager = None + } else { + this.kafkaLinkingManager = Some(kafkaLinkingManager) + } + } + + def setTrafficInterceptor(trafficInterceptor: TrafficInterceptor): Unit = { + this.trafficInterceptor = trafficInterceptor + } + + def setS3StreamContext(ctx: com.automq.stream.Context): Unit = { + this.partitionSnapshotsManager = new PartitionSnapshotsManager(time, config.automq, ctx.confirmWAL(), () => metadataCache.autoMQVersion()) + } + } diff --git a/core/src/main/scala/kafka/server/streamaspect/PartitionLifecycleListener.java b/core/src/main/scala/kafka/server/streamaspect/PartitionLifecycleListener.java new file mode 100644 index 0000000000..6d88bcc1cb --- /dev/null +++ b/core/src/main/scala/kafka/server/streamaspect/PartitionLifecycleListener.java @@ -0,0 +1,30 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.server.streamaspect; + +import kafka.cluster.Partition; + +public interface PartitionLifecycleListener { + + void onOpen(Partition partition); + + void onClose(Partition partition); + +} diff --git a/core/src/test/java/kafka/autobalancer/AutoBalancerListenerTest.java b/core/src/test/java/kafka/autobalancer/AutoBalancerListenerTest.java index 0d24631119..3edd5b9a82 100644 --- a/core/src/test/java/kafka/autobalancer/AutoBalancerListenerTest.java +++ b/core/src/test/java/kafka/autobalancer/AutoBalancerListenerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer; diff --git a/core/src/test/java/kafka/autobalancer/ControllerActionExecutorServiceTest.java b/core/src/test/java/kafka/autobalancer/ControllerActionExecutorServiceTest.java index 91e2f501a1..0447a1ce2b 100644 --- a/core/src/test/java/kafka/autobalancer/ControllerActionExecutorServiceTest.java +++ b/core/src/test/java/kafka/autobalancer/ControllerActionExecutorServiceTest.java @@ -31,12 +31,14 @@ import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import java.util.List; import java.util.concurrent.CompletableFuture; +@Timeout(60) @Tag("S3Unit") public class ControllerActionExecutorServiceTest { diff --git a/core/src/test/java/kafka/autobalancer/LoadRetrieverTest.java b/core/src/test/java/kafka/autobalancer/LoadRetrieverTest.java new file mode 100644 index 0000000000..c10768f374 --- /dev/null +++ b/core/src/test/java/kafka/autobalancer/LoadRetrieverTest.java @@ -0,0 +1,81 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.autobalancer; + +import kafka.autobalancer.config.AutoBalancerControllerConfig; +import kafka.autobalancer.model.ClusterModel; + +import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord; +import org.apache.kafka.common.metadata.RegisterBrokerRecord; +import org.apache.kafka.common.metadata.UnregisterBrokerRecord; +import org.apache.kafka.controller.Controller; +import org.apache.kafka.metadata.BrokerRegistrationFencingChange; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.List; + +public class LoadRetrieverTest { + + @Test + public void testBrokerChanged() { + LoadRetriever loadRetriever = Mockito.spy(new LoadRetriever(Mockito.mock(AutoBalancerControllerConfig.class), Mockito.mock(Controller.class), Mockito.mock(ClusterModel.class))); + loadRetriever.onBrokerRegister(new RegisterBrokerRecord().setBrokerId(0).setFenced(false).setEndPoints( + new RegisterBrokerRecord.BrokerEndpointCollection(List.of( + new RegisterBrokerRecord.BrokerEndpoint().setHost("192.168.0.0").setPort(9092)).iterator()))); + loadRetriever.onBrokerRegister(new RegisterBrokerRecord().setBrokerId(1).setFenced(false).setEndPoints( + new RegisterBrokerRecord.BrokerEndpointCollection(List.of( + new RegisterBrokerRecord.BrokerEndpoint().setHost("192.168.0.1").setPort(9093)).iterator()))); + loadRetriever.checkAndCreateConsumer(0); + + Assertions.assertEquals(loadRetriever.buildBootstrapServer(), "192.168.0.1:9093,192.168.0.0:9092"); + Assertions.assertTrue(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertTrue(loadRetriever.hasAvailableBroker()); + + loadRetriever.onBrokerRegistrationChanged(new BrokerRegistrationChangeRecord().setBrokerId(0).setFenced(BrokerRegistrationFencingChange.FENCE.value())); + Assertions.assertTrue(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertTrue(loadRetriever.hasAvailableBroker()); + + loadRetriever.onBrokerRegistrationChanged(new BrokerRegistrationChangeRecord().setBrokerId(1).setFenced(BrokerRegistrationFencingChange.FENCE.value())); + Assertions.assertFalse(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertFalse(loadRetriever.hasAvailableBroker()); + + loadRetriever.onBrokerRegistrationChanged(new BrokerRegistrationChangeRecord().setBrokerId(1).setFenced(BrokerRegistrationFencingChange.UNFENCE.value())); + Assertions.assertEquals(loadRetriever.buildBootstrapServer(), "192.168.0.1:9093"); + Assertions.assertTrue(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertTrue(loadRetriever.hasAvailableBroker()); + + loadRetriever.onBrokerRegister(new RegisterBrokerRecord().setBrokerId(1).setFenced(false).setEndPoints( + new RegisterBrokerRecord.BrokerEndpointCollection(List.of( + new RegisterBrokerRecord.BrokerEndpoint().setHost("192.168.0.2").setPort(9094)).iterator()))); + Assertions.assertFalse(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertTrue(loadRetriever.hasAvailableBroker()); + + Assertions.assertEquals(loadRetriever.buildBootstrapServer(), "192.168.0.2:9094"); + Assertions.assertTrue(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertTrue(loadRetriever.hasAvailableBroker()); + + loadRetriever.onBrokerUnregister(new UnregisterBrokerRecord().setBrokerId(1)); + Assertions.assertFalse(loadRetriever.hasAvailableBrokerInUse()); + Assertions.assertFalse(loadRetriever.hasAvailableBroker()); + } +} diff --git a/core/src/test/java/kafka/autobalancer/common/normalizer/NormalizerTest.java b/core/src/test/java/kafka/autobalancer/common/normalizer/NormalizerTest.java index 3eeaa2c4a2..386ef61d9d 100644 --- a/core/src/test/java/kafka/autobalancer/common/normalizer/NormalizerTest.java +++ b/core/src/test/java/kafka/autobalancer/common/normalizer/NormalizerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.common.normalizer; diff --git a/core/src/test/java/kafka/autobalancer/config/AutoBalancerConfigTest.java b/core/src/test/java/kafka/autobalancer/config/AutoBalancerConfigTest.java index fd242a5eb5..6f34dcd181 100644 --- a/core/src/test/java/kafka/autobalancer/config/AutoBalancerConfigTest.java +++ b/core/src/test/java/kafka/autobalancer/config/AutoBalancerConfigTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.config; diff --git a/core/src/test/java/kafka/autobalancer/goals/AbstractResourceGoalTest.java b/core/src/test/java/kafka/autobalancer/goals/AbstractResourceGoalTest.java index 88bc22c94c..bd45c80fcc 100644 --- a/core/src/test/java/kafka/autobalancer/goals/AbstractResourceGoalTest.java +++ b/core/src/test/java/kafka/autobalancer/goals/AbstractResourceGoalTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.goals; diff --git a/core/src/test/java/kafka/autobalancer/goals/GoalTestBase.java b/core/src/test/java/kafka/autobalancer/goals/GoalTestBase.java index 71151413d9..90a6e73d9f 100644 --- a/core/src/test/java/kafka/autobalancer/goals/GoalTestBase.java +++ b/core/src/test/java/kafka/autobalancer/goals/GoalTestBase.java @@ -26,6 +26,7 @@ import org.apache.kafka.common.TopicPartition; import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Timeout; import java.util.Collection; import java.util.HashMap; @@ -36,6 +37,7 @@ import static kafka.autobalancer.common.types.Resource.NW_IN; import static kafka.autobalancer.common.types.Resource.NW_OUT; +@Timeout(60) @Tag("S3Unit") public class GoalTestBase { private final Map goalMap = new HashMap<>(); diff --git a/core/src/test/java/kafka/autobalancer/goals/ResourceUsageDistributionGoalTest.java b/core/src/test/java/kafka/autobalancer/goals/ResourceUsageDistributionGoalTest.java index 804fae6432..b2dcb7bce7 100644 --- a/core/src/test/java/kafka/autobalancer/goals/ResourceUsageDistributionGoalTest.java +++ b/core/src/test/java/kafka/autobalancer/goals/ResourceUsageDistributionGoalTest.java @@ -32,6 +32,7 @@ import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -45,6 +46,7 @@ import static kafka.autobalancer.common.types.Resource.NW_IN; import static kafka.autobalancer.common.types.Resource.NW_OUT; +@Timeout(60) @Tag("S3Unit") public class ResourceUsageDistributionGoalTest extends GoalTestBase { diff --git a/core/src/test/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporterTest.java b/core/src/test/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporterTest.java index cbc54359a5..2de7c67498 100644 --- a/core/src/test/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporterTest.java +++ b/core/src/test/java/kafka/autobalancer/metricsreporter/AutoBalancerMetricsReporterTest.java @@ -25,11 +25,13 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.mockito.Mockito; import java.util.HashMap; import java.util.Map; +@Timeout(60) @Tag("S3Unit") public class AutoBalancerMetricsReporterTest { diff --git a/core/src/test/java/kafka/autobalancer/metricsreporter/metric/DerivatorTest.java b/core/src/test/java/kafka/autobalancer/metricsreporter/metric/DerivatorTest.java index 88ceff5cd2..9d0dbf3242 100644 --- a/core/src/test/java/kafka/autobalancer/metricsreporter/metric/DerivatorTest.java +++ b/core/src/test/java/kafka/autobalancer/metricsreporter/metric/DerivatorTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.metricsreporter.metric; diff --git a/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricSerdeTest.java b/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricSerdeTest.java index 1ad52c17f8..8062092684 100644 --- a/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricSerdeTest.java +++ b/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricSerdeTest.java @@ -23,6 +23,7 @@ import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.Map; @@ -30,6 +31,7 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; +@Timeout(60) @Tag("S3Unit") public class MetricSerdeTest { private static final long TIME = 123L; diff --git a/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricsUtilsTest.java b/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricsUtilsTest.java index e0ae29f854..e53eeb65d0 100644 --- a/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricsUtilsTest.java +++ b/core/src/test/java/kafka/autobalancer/metricsreporter/metric/MetricsUtilsTest.java @@ -26,10 +26,12 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.Collections; import java.util.Map; +@Timeout(60) @Tag("S3Unit") public class MetricsUtilsTest { diff --git a/core/src/test/java/kafka/autobalancer/model/ClusterModelTest.java b/core/src/test/java/kafka/autobalancer/model/ClusterModelTest.java index 224a45ab24..e9321bc5a9 100644 --- a/core/src/test/java/kafka/autobalancer/model/ClusterModelTest.java +++ b/core/src/test/java/kafka/autobalancer/model/ClusterModelTest.java @@ -37,12 +37,14 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.Collection; import java.util.Collections; import java.util.Map; import java.util.Set; +@Timeout(60) @Tag("S3Unit") public class ClusterModelTest { @@ -487,20 +489,20 @@ public void testSlowBroker() { clusterModel.unregisterBroker(0); // test high append latency - clusterModel.onBrokerRegister(registerBrokerRecord0); - for (int i = 0; i < 100; i++) { - Assertions.assertTrue(clusterModel.updateBrokerMetrics(0, createBrokerMetrics(0, - 0, 0, 0, MetricVersion.V1).getMetricValueMap().entrySet(), System.currentTimeMillis())); - } - snapshot = clusterModel.snapshot(); - snapshot.markSlowBrokers(); - Assertions.assertFalse(snapshot.broker(0).isSlowBroker()); - Assertions.assertTrue(clusterModel.updateBrokerMetrics(0, createBrokerMetrics(0, - 2000, 0, 0, MetricVersion.V1).getMetricValueMap().entrySet(), System.currentTimeMillis())); - snapshot = clusterModel.snapshot(); - snapshot.markSlowBrokers(); - Assertions.assertTrue(snapshot.broker(0).isSlowBroker()); - clusterModel.unregisterBroker(0); +// clusterModel.onBrokerRegister(registerBrokerRecord0); +// for (int i = 0; i < 100; i++) { +// Assertions.assertTrue(clusterModel.updateBrokerMetrics(0, createBrokerMetrics(0, +// 0, 0, 0, MetricVersion.V1).getMetricValueMap().entrySet(), System.currentTimeMillis())); +// } +// snapshot = clusterModel.snapshot(); +// snapshot.markSlowBrokers(); +// Assertions.assertFalse(snapshot.broker(0).isSlowBroker()); +// Assertions.assertTrue(clusterModel.updateBrokerMetrics(0, createBrokerMetrics(0, +// 2000, 0, 0, MetricVersion.V1).getMetricValueMap().entrySet(), System.currentTimeMillis())); +// snapshot = clusterModel.snapshot(); +// snapshot.markSlowBrokers(); +// Assertions.assertTrue(snapshot.broker(0).isSlowBroker()); +// clusterModel.unregisterBroker(0); // test high pending append latency clusterModel.onBrokerRegister(registerBrokerRecord0); diff --git a/core/src/test/java/kafka/autobalancer/model/samples/SnapshottableSamplesTest.java b/core/src/test/java/kafka/autobalancer/model/samples/SnapshottableSamplesTest.java index 049c46a0c3..7e393bb433 100644 --- a/core/src/test/java/kafka/autobalancer/model/samples/SnapshottableSamplesTest.java +++ b/core/src/test/java/kafka/autobalancer/model/samples/SnapshottableSamplesTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.model.samples; diff --git a/core/src/test/java/kafka/autobalancer/services/AbstractResumableServiceTest.java b/core/src/test/java/kafka/autobalancer/services/AbstractResumableServiceTest.java index c9b92e3f42..2928432a50 100644 --- a/core/src/test/java/kafka/autobalancer/services/AbstractResumableServiceTest.java +++ b/core/src/test/java/kafka/autobalancer/services/AbstractResumableServiceTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.autobalancer.services; diff --git a/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java b/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java new file mode 100644 index 0000000000..f2f426aa84 --- /dev/null +++ b/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java @@ -0,0 +1,195 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.backpressure; + +import kafka.automq.AutoMQConfig; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; + +public class DefaultBackPressureManagerTest { + static String sourceA = "sourceA"; + static String sourceB = "sourceB"; + static String sourceC = "sourceC"; + + BackPressureConfig config; + DefaultBackPressureManager manager; + + Regulator regulator; + int regulatorIncreaseCalled = 0; + int regulatorDecreaseCalled = 0; + + ScheduledExecutorService scheduler; + int schedulerScheduleCalled = 0; + long schedulerScheduleDelay = 0; + + @BeforeEach + public void setup() { + regulator = mock(Regulator.class); + scheduler = mock(ScheduledExecutorService.class); + + // Mock the regulator to count the number of times each method is called + doAnswer(invocation -> { + regulatorIncreaseCalled++; + return null; + }).when(regulator).increase(); + doAnswer(invocation -> { + regulatorDecreaseCalled++; + return null; + }).when(regulator).decrease(); + + // Mock the scheduler to run the scheduled task immediately and only once + doAnswer(invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + return null; + }).when(scheduler).scheduleWithFixedDelay(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)); + doAnswer(invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + schedulerScheduleCalled++; + schedulerScheduleDelay = invocation.getArgument(1); + return null; + }).when(scheduler).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + } + + @Test + public void testDynamicConfig() { + initManager(false, 0); + + callChecker(sourceC, LoadLevel.NORMAL); + callChecker(sourceB, LoadLevel.HIGH); + assertRegulatorCalled(0, 0); + + manager.reconfigure(Map.of( + AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, "true" + )); + callChecker(sourceC, LoadLevel.NORMAL); + callChecker(sourceB, LoadLevel.NORMAL); + assertRegulatorCalled(1, 1); + + manager.reconfigure(Map.of( + AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, "false" + )); + callChecker(sourceC, LoadLevel.NORMAL); + callChecker(sourceB, LoadLevel.HIGH); + assertRegulatorCalled(1, 1); + } + + @Test + public void testPriority1() { + initManager(0); + + callChecker(sourceB, LoadLevel.HIGH); + callChecker(sourceC, LoadLevel.NORMAL); + + assertRegulatorCalled(0, 2); + } + + @Test + public void testPriority2() { + initManager(0); + + callChecker(sourceC, LoadLevel.NORMAL); + callChecker(sourceB, LoadLevel.HIGH); + + assertRegulatorCalled(1, 1); + } + + @Test + public void testOverride() { + initManager(0); + + callChecker(sourceA, LoadLevel.NORMAL); + callChecker(sourceA, LoadLevel.HIGH); + callChecker(sourceA, LoadLevel.NORMAL); + + assertRegulatorCalled(2, 1); + } + + @Test + public void testCooldown() { + final long cooldownMs = Long.MAX_VALUE; + final long tolerance = 1000; + + initManager(cooldownMs); + + callChecker(sourceA, LoadLevel.HIGH); + assertRegulatorCalled(0, 0); + assertSchedulerCalled(1); + assertEquals(cooldownMs, schedulerScheduleDelay, tolerance); + + callChecker(sourceA, LoadLevel.NORMAL); + assertRegulatorCalled(0, 0); + assertSchedulerCalled(2); + assertEquals(cooldownMs, schedulerScheduleDelay, tolerance); + } + + private void initManager(long cooldownMs) { + initManager(true, cooldownMs); + } + + /** + * Should be called at the beginning of each test to initialize the manager. + */ + private void initManager(boolean enabled, long cooldownMs) { + config = new BackPressureConfig(enabled, cooldownMs); + manager = new DefaultBackPressureManager(config, regulator); + manager.checkerScheduler = scheduler; + } + + private void callChecker(String source, LoadLevel level) { + manager.registerChecker(new Checker() { + @Override + public String source() { + return source; + } + + @Override + public LoadLevel check() { + return level; + } + + @Override + public long intervalMs() { + return 1; + } + }); + } + + private void assertRegulatorCalled(int increase, int decrease) { + assertEquals(increase, regulatorIncreaseCalled); + assertEquals(decrease, regulatorDecreaseCalled); + } + + private void assertSchedulerCalled(int times) { + assertEquals(times, schedulerScheduleCalled); + } +} diff --git a/core/src/test/java/kafka/automq/failover/FailoverContextTest.java b/core/src/test/java/kafka/automq/failover/FailoverContextTest.java new file mode 100644 index 0000000000..e3e60503bc --- /dev/null +++ b/core/src/test/java/kafka/automq/failover/FailoverContextTest.java @@ -0,0 +1,44 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.failover; + +import kafka.automq.utils.JsonUtils; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +@Tag("S3Unit") +public class FailoverContextTest { + + @Test + public void testEncodeDecodeV1() { + FailoverContext context = new FailoverContext(111, 333, 222, "kraftWalConfigs"); + String encoded = JsonUtils.encode(context); + FailoverContext decoded = JsonUtils.decode(encoded, FailoverContext.class); + assertEquals("{\"n\":111,\"t\":222,\"e\":333,\"c\":\"kraftWalConfigs\"}", encoded); + assertEquals(context.getNodeId(), decoded.getNodeId()); + assertEquals(context.getNodeEpoch(), decoded.getNodeEpoch()); + assertEquals(context.getTarget(), decoded.getTarget()); + assertEquals(context.getKraftWalConfigs(), decoded.getKraftWalConfigs()); + } + +} diff --git a/core/src/test/java/kafka/automq/partition/snapshot/ConfirmWalDataDeltaTest.java b/core/src/test/java/kafka/automq/partition/snapshot/ConfirmWalDataDeltaTest.java new file mode 100644 index 0000000000..939c1af946 --- /dev/null +++ b/core/src/test/java/kafka/automq/partition/snapshot/ConfirmWalDataDeltaTest.java @@ -0,0 +1,149 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.partition.snapshot; + +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; + +import com.automq.stream.s3.ConfirmWAL; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import io.netty.buffer.Unpooled; + +import static kafka.automq.partition.snapshot.ConfirmWalDataDelta.MAX_RECORDS_BUFFER_SIZE; +import static kafka.automq.partition.snapshot.ConfirmWalDataDelta.STATE_NOT_SYNC; +import static kafka.automq.partition.snapshot.ConfirmWalDataDelta.STATE_SYNCING; +import static kafka.automq.partition.snapshot.ConfirmWalDataDelta.decodeDeltaRecords; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ConfirmWalDataDeltaTest { + + ConfirmWAL confirmWAL; + ConfirmWalDataDelta delta; + long walOffset = 233; + long nextWalOffset = walOffset; + + @BeforeEach + void setup() { + confirmWAL = mock(ConfirmWAL.class); + when(confirmWAL.addAppendListener(any())).thenReturn(() -> { + }); + delta = new ConfirmWalDataDelta(confirmWAL); + walOffset = 233; + nextWalOffset = walOffset; + } + + @Test + public void testHandle() { + AutomqGetPartitionSnapshotResponseData resp = new AutomqGetPartitionSnapshotResponseData(); + // There's no new wal data. + DefaultRecordOffset confirmOffset = DefaultRecordOffset.of(1, 233, 0); + when(confirmWAL.confirmOffset()).thenReturn(confirmOffset); + delta.handle((short) 1, resp); + assertEquals(confirmOffset, DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset()))); + // In requestVersion=1, the confirmWalDeltaData is an empty array. + assertEquals(0, resp.confirmWalDeltaData().length); + + // In requestVersion=2, the confirmWalDeltaData is null when there isn't new wal data, or it's not in STATE_SYNCING. + resp = new AutomqGetPartitionSnapshotResponseData(); + when(confirmWAL.confirmOffset()).thenReturn(confirmOffset); + delta.handle((short) 2, resp); + assertEquals(confirmOffset, DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset()))); + Assertions.assertNull(resp.confirmWalDeltaData()); + assertEquals(STATE_NOT_SYNC, delta.state); + + // New record has been appended, the state will change from STATE_NOT_SYNC to STATE_SYNCING. + when(confirmWAL.confirmOffset()).thenThrow(new UnsupportedOperationException()); + for (int i = 0; i < 64; i++) { + for (int j = 0; j < 3; j++) { + onAppend(3 * i + j); + } + resp = new AutomqGetPartitionSnapshotResponseData(); + delta.handle((short) 2, resp); + assertEquals(DefaultRecordOffset.of(1, nextWalOffset, 0), DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset()))); + if (i == 0) { + // The first response in STATE_SYNCING only take confirmOffset from wal records and set the confirmWalDeltaData null. + Assertions.assertNull(resp.confirmWalDeltaData()); + } else { + List recordList = decodeDeltaRecords(resp.confirmWalDeltaData()); + assertEquals(3, recordList.size()); + for (int j = 0; j < 3; j++) { + StreamRecordBatch record = recordList.get(j); + assertEquals(3 * i + j, record.getBaseOffset()); + assertEquals(1024, record.getPayload().readableBytes()); + record.release(); + } + } + assertEquals(0, delta.size.get()); + assertEquals(STATE_SYNCING, delta.state); + } + } + + @Test + public void testOnAppend_bufferExceed() { + AutomqGetPartitionSnapshotResponseData resp = new AutomqGetPartitionSnapshotResponseData(); + + onAppend(3); + + resp = new AutomqGetPartitionSnapshotResponseData(); + delta.handle((short) 2, resp); + assertEquals(nextWalOffset, DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset())).offset()); + Assertions.assertNull(resp.confirmWalDeltaData()); + assertEquals(STATE_SYNCING, delta.state); + + // buffer exceed + int i = 0; + for (int size = 0; size < MAX_RECORDS_BUFFER_SIZE; i++) { + size += onAppend(4 + i); + } + assertEquals(0, delta.size.get()); + assertEquals(STATE_NOT_SYNC, delta.state); + + onAppend(4 + i); + resp = new AutomqGetPartitionSnapshotResponseData(); + delta.handle((short) 2, resp); + assertEquals(nextWalOffset, DefaultRecordOffset.of(Unpooled.wrappedBuffer(resp.confirmWalEndOffset())).offset()); + Assertions.assertNull(resp.confirmWalDeltaData()); + assertEquals(STATE_SYNCING, delta.state); + assertEquals(0, delta.size.get()); + } + + int onAppend(long recordBaseOffset) { + StreamRecordBatch record = StreamRecordBatch.of(1, 2, recordBaseOffset, 1, Unpooled.wrappedBuffer(new byte[1024])); + nextWalOffset = walOffset + record.encoded().readableBytes(); + delta.onAppend( + record, + DefaultRecordOffset.of(1, walOffset, record.encoded().readableBytes()), + DefaultRecordOffset.of(1, nextWalOffset, 0) + ); + walOffset = nextWalOffset; + return record.encoded().readableBytes(); + } + +} diff --git a/core/src/test/java/kafka/automq/table/CatalogFactoryTest.java b/core/src/test/java/kafka/automq/table/CatalogFactoryTest.java new file mode 100644 index 0000000000..9ef3e531b0 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/CatalogFactoryTest.java @@ -0,0 +1,146 @@ +package kafka.automq.table; + +import kafka.server.KafkaConfig; + +import org.apache.kafka.raft.QuorumConfig; +import org.apache.kafka.server.config.KRaftConfigs; + +import com.sun.net.httpserver.HttpServer; + +import org.apache.iceberg.aws.s3.S3FileIOProperties; +import org.apache.iceberg.inmemory.InMemoryFileIO; +import org.apache.iceberg.rest.RESTCatalog; +import org.apache.iceberg.util.SerializableMap; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.stream.Stream; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.stream.Collectors.toMap; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; + +@TestInstance(PER_CLASS) +class CatalogFactoryTest { + // minimalistic properties to let KafkaConfig validation pass and let us test our catalog factory + private final Map requiredKafkaConfigProperties = Map.of( + KRaftConfigs.CONTROLLER_LISTENER_NAMES_CONFIG, "PLAINTEXT", + KRaftConfigs.NODE_ID_CONFIG, "2", + KRaftConfigs.PROCESS_ROLES_CONFIG, "controller", + QuorumConfig.QUORUM_VOTERS_CONFIG, "2@localhost:9092" + ); + + @Test + void restPassthroughProperties() throws IOException { + final var restCatalog = new RestCatalogMock(); + try (final var autoClose = restCatalog) { + final var config = new KafkaConfig(merge(requiredKafkaConfigProperties, Map.of( + "automq.table.topic.catalog.type", "rest", + "automq.table.topic.catalog.uri", restCatalog.base(), + "automq.table.topic.catalog.header.x-custom", "my-x", // Apache Polaris needs a tenant header for ex + // automq specific/enforced (not standard catalog passthrough) + "s3.data.buckets", "0@s3://my_bucket?region=us-east-1&endpoint=http://localhost:12345&pathStyle=true" + ))); + final var catalog = new CatalogFactory.Builder(config).build(); + assertInstanceOf(RESTCatalog.class, catalog).close(); + } + assertEquals(List.of("GET /v1/config?warehouse=s3://my_bucket/iceberg\nmy-x"), restCatalog.requests()); + } + + @Test + void ignoreEmptyS3EndpointForRestCatalog() throws IOException { + FakeS3IO.lastS3FileIOProperties = null; + try (final var restCatalog = new RestCatalogMock()) { + final var config = new KafkaConfig(merge(requiredKafkaConfigProperties, Map.of( + "automq.table.topic.catalog.type", "rest", + "automq.table.topic.catalog.uri", restCatalog.base(), + "automq.table.topic.catalog.io-impl", FakeS3IO.class.getName(), + "s3.data.buckets", "0@s3://my_bucket?region=us-east-1" + ))); + final var catalog = new CatalogFactory.Builder(config).build(); + assertInstanceOf(RESTCatalog.class, catalog).close(); + assertNull(FakeS3IO.lastS3FileIOProperties().endpoint(), "S3FileIO endpoint should be null when not set - not even empty"); + } finally { + FakeS3IO.lastS3FileIOProperties = null; + } + } + + @SafeVarargs + private Map merge(final Map... all) { + return Stream.of(all) + .flatMap(it -> it.entrySet().stream()) + .collect(toMap(Map.Entry::getKey, Map.Entry::getValue, (a, b) -> b)); + } + + public static class FakeS3IO extends InMemoryFileIO { + private static S3FileIOProperties lastS3FileIOProperties; + + private static S3FileIOProperties lastS3FileIOProperties() { + return lastS3FileIOProperties; + } + + @Override + public void initialize(final Map properties) { + lastS3FileIOProperties = new S3FileIOProperties(SerializableMap.copyOf(properties)); + super.initialize(properties); + } + } + + private static class RestCatalogMock implements AutoCloseable { + private final List requests = new CopyOnWriteArrayList<>(); // normally overkill but makes the test more accurate + private final HttpServer catalogBackend; + + private RestCatalogMock() throws IOException { + catalogBackend = HttpServer.create(new InetSocketAddress("localhost", 0), 16); + catalogBackend.createContext("/").setHandler(ex -> { + try (ex) { + final var method = ex.getRequestMethod(); + requests.add( + method + ' ' + ex.getRequestURI().getPath() + '?' + ex.getRequestURI().getQuery() + + ('\n' + String.join("", ex.getRequestHeaders().getOrDefault("x-custom", List.of()))) + + ('\n' + new String(ex.getRequestBody().readAllBytes(), UTF_8)).strip()); + + if (method.equals("GET") && + ex.getRequestURI().getPath().equals("/v1/config") && + "warehouse=s3%3A%2F%2Fmy_bucket%2Ficeberg".equals(ex.getRequestURI().getRawQuery())) { + final var body = """ + { + "defaults": {}, + "overrides": {} + } + """.getBytes(UTF_8); + ex.getResponseHeaders().add("content-type", "application/json"); + ex.sendResponseHeaders(200, body.length); + ex.getResponseBody().write(body); + return; + } + + // else we just called an unexpected endpoint, issue a HTTP 404 + ex.sendResponseHeaders(404, 0); + } + }); + catalogBackend.start(); + } + + private String base() { + return "http://localhost:" + catalogBackend.getAddress().getPort(); + } + + private List requests() { + return requests; + } + + @Override + public void close() { + catalogBackend.stop(0); + } + } +} diff --git a/core/src/test/java/kafka/automq/table/binder/AvroRecordBinderTest.java b/core/src/test/java/kafka/automq/table/binder/AvroRecordBinderTest.java new file mode 100644 index 0000000000..cfffed5485 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/binder/AvroRecordBinderTest.java @@ -0,0 +1,623 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.binder; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.util.Utf8; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.avro.CodecSetup; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@Tag("S3Unit") +class AvroRecordBinderTest { + + private static final String TEST_NAMESPACE = "kafka.automq.table.binder"; + + static { + CodecSetup.setup(); + } + + /** + * Tests that when the same Schema instance is used in multiple places (direct field and list element), + * the RecordBinder correctly shares the same binder for that schema instance. + * This verifies the IdentityHashMap optimization. + */ + @Test + public void testStructSchemaInstanceReuseSharesBinder() { + Schema sharedStruct = Schema.createRecord("SharedStruct", null, TEST_NAMESPACE, false); + sharedStruct.setFields(Arrays.asList( + new Schema.Field("value", Schema.create(Schema.Type.LONG), null, null) + )); + + Schema listSchema = Schema.createArray(sharedStruct); + + Schema parent = Schema.createRecord("SharedStructReuseRoot", null, TEST_NAMESPACE, false); + parent.setFields(Arrays.asList( + new Schema.Field("directField", sharedStruct, null, null), + new Schema.Field("listField", listSchema, null, null) + )); + + GenericRecord directValue = new GenericData.Record(sharedStruct); + directValue.put("value", 1L); + + @SuppressWarnings("unchecked") + GenericData.Array listValue = new GenericData.Array<>(2, listSchema); + GenericRecord listEntry1 = new GenericData.Record(sharedStruct); + listEntry1.put("value", 2L); + listValue.add(listEntry1); + GenericRecord listEntry2 = new GenericData.Record(sharedStruct); + listEntry2.put("value", 3L); + listValue.add(listEntry2); + + GenericRecord parentRecord = new GenericData.Record(parent); + parentRecord.put("directField", directValue); + parentRecord.put("listField", listValue); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(parent); + Record icebergRecord = new RecordBinder(icebergSchema, parent).bind(parentRecord); + + Record directRecord = (Record) icebergRecord.getField("directField"); + assertEquals(1L, directRecord.getField("value")); + + @SuppressWarnings("unchecked") + List boundList = (List) icebergRecord.getField("listField"); + assertEquals(2, boundList.size()); + assertEquals(2L, boundList.get(0).getField("value")); + assertEquals(3L, boundList.get(1).getField("value")); + } + + /** + * Tests that structs with the same full name but different schemas in different contexts + * (direct field vs list element) are handled correctly using IdentityHashMap. + * This ensures schema identity, not name equality, is used for binder lookup. + */ + @Test + public void testStructBindersHandleDuplicateFullNames() { + Schema directStruct = Schema.createRecord("DuplicatedStruct", null, TEST_NAMESPACE, false); + directStruct.setFields(Arrays.asList( + new Schema.Field("directOnly", Schema.create(Schema.Type.STRING), null, null) + )); + + Schema listStruct = Schema.createRecord("DuplicatedStruct", null, TEST_NAMESPACE, false); + listStruct.setFields(Arrays.asList( + new Schema.Field("listOnly", Schema.create(Schema.Type.INT), null, null) + )); + + Schema listSchema = Schema.createArray(listStruct); + + Schema parent = Schema.createRecord("StructCollisionRoot", null, TEST_NAMESPACE, false); + parent.setFields(Arrays.asList( + new Schema.Field("directField", directStruct, null, null), + new Schema.Field("listField", listSchema, null, null) + )); + + GenericRecord parentRecord = new GenericData.Record(parent); + GenericRecord directRecord = new GenericData.Record(directStruct); + directRecord.put("directOnly", new Utf8("direct")); + parentRecord.put("directField", directRecord); + + @SuppressWarnings("unchecked") + GenericData.Array listValue = new GenericData.Array<>(1, listSchema); + GenericRecord listRecord = new GenericData.Record(listStruct); + listRecord.put("listOnly", 42); + listValue.add(listRecord); + parentRecord.put("listField", listValue); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(parent); + Record icebergRecord = new RecordBinder(icebergSchema, parent).bind(parentRecord); + + Record directField = (Record) icebergRecord.getField("directField"); + assertEquals("direct", directField.getField("directOnly").toString()); + + @SuppressWarnings("unchecked") + List boundList = (List) icebergRecord.getField("listField"); + assertEquals(1, boundList.size()); + assertEquals(42, boundList.get(0).getField("listOnly")); + } + + /** + * Tests duplicate struct names in map values context. + * Verifies IdentityHashMap correctly distinguishes between schemas with same name. + */ + @Test + public void testStructBindersHandleDuplicateFullNamesInMapValues() { + Schema directStruct = Schema.createRecord("DuplicatedStruct", null, TEST_NAMESPACE, false); + directStruct.setFields(Arrays.asList( + new Schema.Field("directOnly", Schema.create(Schema.Type.STRING), null, null) + )); + + Schema mapStruct = Schema.createRecord("DuplicatedStruct", null, TEST_NAMESPACE, false); + mapStruct.setFields(Arrays.asList( + new Schema.Field("mapOnly", Schema.create(Schema.Type.LONG), null, null) + )); + + Schema mapSchema = Schema.createMap(mapStruct); + + Schema parent = Schema.createRecord("StructCollisionMapRoot", null, TEST_NAMESPACE, false); + parent.setFields(Arrays.asList( + new Schema.Field("directField", directStruct, null, null), + new Schema.Field("mapField", mapSchema, null, null) + )); + + GenericRecord parentRecord = new GenericData.Record(parent); + GenericRecord directRecord = new GenericData.Record(directStruct); + directRecord.put("directOnly", new Utf8("direct")); + parentRecord.put("directField", directRecord); + + Map mapValue = new HashMap<>(); + GenericRecord mapEntry = new GenericData.Record(mapStruct); + mapEntry.put("mapOnly", 123L); + mapValue.put("key", mapEntry); + parentRecord.put("mapField", mapValue); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(parent); + Record icebergRecord = new RecordBinder(icebergSchema, parent).bind(parentRecord); + + Record directField = (Record) icebergRecord.getField("directField"); + assertEquals("direct", directField.getField("directOnly").toString()); + + @SuppressWarnings("unchecked") + Map boundMap = (Map) icebergRecord.getField("mapField"); + assertEquals(1, boundMap.size()); + assertEquals(123L, boundMap.get("key").getField("mapOnly")); + } + + /** + * Tests that AvroValueAdapter throws IllegalStateException when trying to convert + * a struct with missing fields in the source Avro schema. + */ + @Test + public void testConvertStructThrowsWhenSourceFieldMissing() { + Schema nestedSchema = Schema.createRecord("NestedRecord", null, TEST_NAMESPACE, false); + nestedSchema.setFields(Arrays.asList( + new Schema.Field("presentField", Schema.create(Schema.Type.STRING), null, null) + )); + + GenericRecord nestedRecord = new GenericData.Record(nestedSchema); + nestedRecord.put("presentField", new Utf8("value")); + + Types.StructType icebergStruct = Types.StructType.of( + Types.NestedField.optional(2, "presentField", Types.StringType.get()), + Types.NestedField.optional(3, "missingField", Types.StringType.get()) + ); + + AvroValueAdapter adapter = new AvroValueAdapter(); + IllegalStateException exception = assertThrows(IllegalStateException.class, + () -> adapter.convert(nestedRecord, nestedSchema, icebergStruct)); + assertTrue(exception.getMessage().contains("missingField")); + assertTrue(exception.getMessage().contains("NestedRecord")); + } + + /** + * Tests field count statistics for various field types and sizes. + * Verifies that small/large strings, binary fields, and primitives are counted correctly. + */ + @Test + public void testFieldCountStatistics() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"TestRecord\",\n" + + " \"fields\": [\n" + + " {\"name\": \"smallString\", \"type\": \"string\"},\n" + + " {\"name\": \"largeString\", \"type\": \"string\"},\n" + + " {\"name\": \"intField\", \"type\": \"int\"},\n" + + " {\"name\": \"binaryField\", \"type\": \"bytes\"},\n" + + " {\"name\": \"optionalStringField\", \"type\": [\"null\", \"string\"], \"default\": null}\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("smallString", "small"); // 5 chars = 3 field + avroRecord.put("largeString", "a".repeat(50)); // 50 chars = 3 + 50/32 = 4 + avroRecord.put("intField", 42); // primitive = 1 field + avroRecord.put("binaryField", ByteBuffer.wrap("test".repeat(10).getBytes())); // 5 + avroRecord.put("optionalStringField", "optional"); + + Record icebergRecord = recordBinder.bind(avroRecord); + + // Access all fields to trigger counting + assertEquals("small", icebergRecord.getField("smallString")); + assertEquals("a".repeat(50), icebergRecord.getField("largeString")); + assertEquals(42, icebergRecord.getField("intField")); + assertEquals("test".repeat(10), new String(((ByteBuffer) icebergRecord.getField("binaryField")).array())); + assertEquals("optional", icebergRecord.getField("optionalStringField").toString()); + + long fieldCount = recordBinder.getAndResetFieldCount(); + assertEquals(16, fieldCount); + + // Second call should return 0 (reset) + assertEquals(0, recordBinder.getAndResetFieldCount()); + } + + /** + * Tests field counting for complex types (LIST and MAP). + * Verifies that list and map elements are counted correctly. + */ + @Test + public void testFieldCountWithComplexTypes() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"ComplexRecord\",\n" + + " \"fields\": [\n" + + " {\"name\": \"stringList\", \"type\": {\"type\": \"array\", \"items\": \"string\"}},\n" + + " {\"name\": \"stringMap\", \"type\": {\"type\": \"map\", \"values\": \"string\"}}\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("stringList", Arrays.asList("a", "b", "c")); + + Map map = new HashMap<>(); + map.put("key1", "val1"); + map.put("key2", "val2"); + avroRecord.put("stringMap", map); + + Record icebergRecord = recordBinder.bind(avroRecord); + + // Access fields to trigger counting + icebergRecord.getField("stringList"); + icebergRecord.getField("stringMap"); + + // Total: 10 (list) + 13 (map) = 23 fields + long fieldCount = recordBinder.getAndResetFieldCount(); + assertEquals(23, fieldCount); + } + + /** + * Tests field counting for nested struct fields. + * Verifies that nested struct fields contribute to the count correctly. + */ + @Test + public void testFieldCountWithNestedStructure() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"NestedRecord\",\n" + + " \"fields\": [\n" + + " {\"name\": \"simpleField\", \"type\": \"string\"},\n" + + " {\n" + + " \"name\": \"nestedField\",\n" + + " \"type\": {\n" + + " \"type\": \"record\",\n" + + " \"name\": \"Nested\",\n" + + " \"fields\": [\n" + + " {\"name\": \"nestedString\", \"type\": \"string\"},\n" + + " {\"name\": \"nestedInt\", \"type\": \"int\"}\n" + + " ]\n" + + " }\n" + + " }\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + GenericRecord nestedRecord = new GenericData.Record(avroSchema.getField("nestedField").schema()); + nestedRecord.put("nestedString", "nested"); + nestedRecord.put("nestedInt", 123); + + GenericRecord mainRecord = new GenericData.Record(avroSchema); + mainRecord.put("simpleField", "simple"); + mainRecord.put("nestedField", nestedRecord); + + Record icebergRecord = recordBinder.bind(mainRecord); + + // Access all fields including nested ones + assertEquals("simple", icebergRecord.getField("simpleField")); + Record nested = (Record) icebergRecord.getField("nestedField"); + assertEquals("nested", nested.getField("nestedString")); + assertEquals(123, nested.getField("nestedInt")); + + // Total: 3 (simple) + 1(struct) + 3 (nested string) + 1 (nested int) = 8 fields + long fieldCount = recordBinder.getAndResetFieldCount(); + assertEquals(8, fieldCount); + } + + /** + * Tests that field counts accumulate across multiple record bindings. + * Verifies batch processing statistics. + */ + @Test + public void testFieldCountBatchAccumulation() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"SimpleRecord\",\n" + + " \"fields\": [\n" + + " {\"name\": \"stringField\", \"type\": \"string\"},\n" + + " {\"name\": \"intField\", \"type\": \"int\"}\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + // Process multiple records + for (int i = 0; i < 3; i++) { + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("stringField", "test" + i); + avroRecord.put("intField", i); + + Record icebergRecord = recordBinder.bind(avroRecord); + // Access fields to trigger counting + icebergRecord.getField("stringField"); + icebergRecord.getField("intField"); + } + + // Total: 3 records * 4 fields each = 12 fields + long totalFieldCount = recordBinder.getAndResetFieldCount(); + assertEquals(12, totalFieldCount); + } + + /** + * Tests that null values don't contribute to field count. + */ + @Test + public void testFieldCountWithNullValues() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"NullableRecord\",\n" + + " \"fields\": [\n" + + " {\"name\": \"nonNullField\", \"type\": \"string\"},\n" + + " {\"name\": \"nullField\", \"type\": [\"null\", \"string\"], \"default\": null}\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("nonNullField", "value"); + avroRecord.put("nullField", null); + + Record icebergRecord = recordBinder.bind(avroRecord); + + // Access both fields + assertEquals("value", icebergRecord.getField("nonNullField")); + assertNull(icebergRecord.getField("nullField")); + + // Only the non-null field should count + long fieldCount = recordBinder.getAndResetFieldCount(); + assertEquals(3, fieldCount); + } + + /** + * Tests field counting for optional union fields with both null and non-null values. + */ + @Test + public void testFieldCountWithUnionFields() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"UnionCountRecord\",\n" + + " \"fields\": [\n" + + " {\"name\": \"optionalString\", \"type\": [\"null\", \"string\"], \"default\": null}\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + // Test with non-null value + GenericRecord nonNullRecord = new GenericData.Record(avroSchema); + nonNullRecord.put("optionalString", "value"); + + Record icebergRecord = recordBinder.bind(nonNullRecord); + assertEquals("value", icebergRecord.getField("optionalString").toString()); + assertEquals(3, recordBinder.getAndResetFieldCount()); + + // Test with null value + GenericRecord nullRecord = new GenericData.Record(avroSchema); + nullRecord.put("optionalString", null); + + Record nullIcebergRecord = recordBinder.bind(nullRecord); + assertNull(nullIcebergRecord.getField("optionalString")); + assertEquals(0, recordBinder.getAndResetFieldCount()); + } + + /** + * Tests that binding a null GenericRecord returns null. + */ + @Test + public void testBindNullRecordReturnsNull() { + Schema avroSchema = Schema.createRecord("TestRecord", null, TEST_NAMESPACE, false); + avroSchema.setFields(Arrays.asList( + new Schema.Field("field", Schema.create(Schema.Type.STRING), null, null) + )); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + + Record result = recordBinder.bind(null); + assertNull(result); + } + + /** + * Tests that accessing a field with negative position throws IndexOutOfBoundsException. + */ + @Test + public void testGetFieldWithNegativePositionThrowsException() { + Schema avroSchema = Schema.createRecord("TestRecord", null, TEST_NAMESPACE, false); + avroSchema.setFields(Arrays.asList( + new Schema.Field("field", Schema.create(Schema.Type.STRING), null, null) + )); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("field", new Utf8("value")); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + Record icebergRecord = recordBinder.bind(avroRecord); + + IndexOutOfBoundsException exception = assertThrows(IndexOutOfBoundsException.class, + () -> icebergRecord.get(-1)); + assertTrue(exception.getMessage().contains("out of bounds")); + } + + /** + * Tests that accessing a field with position >= size throws IndexOutOfBoundsException. + */ + @Test + public void testGetFieldWithExcessivePositionThrowsException() { + Schema avroSchema = Schema.createRecord("TestRecord", null, TEST_NAMESPACE, false); + avroSchema.setFields(Arrays.asList( + new Schema.Field("field", Schema.create(Schema.Type.STRING), null, null) + )); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("field", new Utf8("value")); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + Record icebergRecord = recordBinder.bind(avroRecord); + + IndexOutOfBoundsException exception = assertThrows(IndexOutOfBoundsException.class, + () -> icebergRecord.get(999)); + assertTrue(exception.getMessage().contains("out of bounds")); + } + + /** + * Tests that accessing a field by an unknown name returns null. + */ + @Test + public void testGetFieldByUnknownNameReturnsNull() { + Schema avroSchema = Schema.createRecord("TestRecord", null, TEST_NAMESPACE, false); + avroSchema.setFields(Arrays.asList( + new Schema.Field("existingField", Schema.create(Schema.Type.STRING), null, null) + )); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("existingField", new Utf8("value")); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + Record icebergRecord = recordBinder.bind(avroRecord); + + assertNull(icebergRecord.getField("nonExistentField")); + } + + /** + * Tests that a UNION containing only NULL type throws IllegalArgumentException. + */ + @Test + public void testUnionWithOnlyNullThrowsException() { + Schema nullOnlyUnion = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL))); + + Schema avroSchema = Schema.createRecord("TestRecord", null, TEST_NAMESPACE, false); + avroSchema.setFields(Arrays.asList( + new Schema.Field("nullField", nullOnlyUnion, null, null) + )); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> new RecordBinder(icebergSchema, avroSchema)); + assertTrue(exception.getMessage().contains("UNION schema contains only NULL type")); + } + + /** + * Tests that null elements in Map-as-Array representation are skipped. + */ + @Test + public void testMapAsArrayWithNullElementsSkipped() { + String avroSchemaStr = "{\n" + + " \"type\": \"record\",\n" + + " \"name\": \"MapAsArrayRecord\",\n" + + " \"fields\": [\n" + + " {\n" + + " \"name\": \"mapField\",\n" + + " \"type\": {\n" + + " \"type\": \"array\",\n" + + " \"logicalType\": \"map\",\n" + + " \"items\": {\n" + + " \"type\": \"record\",\n" + + " \"name\": \"MapEntry\",\n" + + " \"fields\": [\n" + + " {\"name\": \"key\", \"type\": \"string\"},\n" + + " {\"name\": \"value\", \"type\": \"int\"}\n" + + " ]\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + "}"; + + Schema avroSchema = new Schema.Parser().parse(avroSchemaStr); + Schema entrySchema = avroSchema.getField("mapField").schema().getElementType(); + + @SuppressWarnings("unchecked") + GenericData.Array arrayValue = new GenericData.Array<>(3, avroSchema.getField("mapField").schema()); + + // Add valid entry + GenericRecord entry1 = new GenericData.Record(entrySchema); + entry1.put("key", new Utf8("key1")); + entry1.put("value", 100); + arrayValue.add(entry1); + + // Add null entry (should be skipped) + arrayValue.add(null); + + // Add another valid entry + GenericRecord entry2 = new GenericData.Record(entrySchema); + entry2.put("key", new Utf8("key2")); + entry2.put("value", 200); + arrayValue.add(entry2); + + GenericRecord avroRecord = new GenericData.Record(avroSchema); + avroRecord.put("mapField", arrayValue); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(avroSchema); + RecordBinder recordBinder = new RecordBinder(icebergSchema, avroSchema); + Record icebergRecord = recordBinder.bind(avroRecord); + + @SuppressWarnings("unchecked") + Map mapField = (Map) icebergRecord.getField("mapField"); + + // Should only contain 2 entries (null entry skipped) + assertEquals(2, mapField.size()); + assertEquals(100, mapField.get(new Utf8("key1"))); + assertEquals(200, mapField.get(new Utf8("key2"))); + } +} diff --git a/core/src/test/java/kafka/automq/table/binder/AvroRecordBinderTypeTest.java b/core/src/test/java/kafka/automq/table/binder/AvroRecordBinderTypeTest.java new file mode 100644 index 0000000000..a8f6d34571 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/binder/AvroRecordBinderTypeTest.java @@ -0,0 +1,1019 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.binder; + +import com.google.common.collect.ImmutableMap; + +import org.apache.avro.Conversions; +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.DatumReader; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.Decoder; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.Encoder; +import org.apache.avro.io.EncoderFactory; +import org.apache.avro.specific.SpecificDatumReader; +import org.apache.avro.specific.SpecificDatumWriter; +import org.apache.avro.util.Utf8; +import org.apache.commons.lang3.RandomStringUtils; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Table; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.avro.CodecSetup; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericAppenderFactory; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.FileAppenderFactory; +import org.apache.iceberg.io.OutputFileFactory; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.io.UnpartitionedWriter; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.DateTimeUtil; +import org.apache.iceberg.util.UUIDUtil; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockitoAnnotations; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.OffsetDateTime; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; + +import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES; +import static org.apache.iceberg.TableProperties.WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; + +public class AvroRecordBinderTypeTest { + + private static final String TEST_NAMESPACE = "kafka.automq.table.binder"; + + private InMemoryCatalog catalog; + private Table table; + private TaskWriter writer; + private int tableCounter; + + static { + CodecSetup.setup(); + } + + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + catalog = new InMemoryCatalog(); + catalog.initialize("test", ImmutableMap.of()); + catalog.createNamespace(Namespace.of("default")); + tableCounter = 0; + } + + // Test method for converting a single string field + @Test + public void testStringConversion() { + assertFieldRoundTrips("String", "stringField", + () -> Schema.create(Schema.Type.STRING), + schema -> "test_string", + value -> assertEquals("test_string", value.toString()) + ); + } + + // Test method for converting a single integer field + @Test + public void testIntegerConversion() { + assertFieldRoundTrips("Int", "intField", + () -> Schema.create(Schema.Type.INT), + schema -> 42, + value -> assertEquals(42, value) + ); + } + + // Test method for converting a single long field + @Test + public void testLongConversion() { + assertFieldRoundTrips("Long", "longField", + () -> Schema.create(Schema.Type.LONG), + schema -> 123456789L, + value -> assertEquals(123456789L, value) + ); + } + + // Test method for converting a single float field + @Test + public void testFloatConversion() { + assertFieldRoundTrips("Float", "floatField", + () -> Schema.create(Schema.Type.FLOAT), + schema -> 3.14f, + value -> assertEquals(3.14f, (Float) value) + ); + } + + // Test method for converting a single double field + @Test + public void testDoubleConversion() { + assertFieldRoundTrips("Double", "doubleField", + () -> Schema.create(Schema.Type.DOUBLE), + schema -> 6.28, + value -> assertEquals(6.28, value) + ); + } + + // Test method for converting a single boolean field + @Test + public void testBooleanConversion() { + assertFieldRoundTrips("Boolean", "booleanField", + () -> Schema.create(Schema.Type.BOOLEAN), + schema -> true, + value -> assertEquals(true, value) + ); + } + + // Test method for converting a single date field (number of days from epoch) + @Test + public void testDateConversion() { + LocalDate localDate = LocalDate.of(2020, 1, 1); + int epochDays = (int) ChronoUnit.DAYS.between(LocalDate.ofEpochDay(0), localDate); + assertFieldRoundTrips("Date", "dateField", + () -> LogicalTypes.date().addToSchema(Schema.create(Schema.Type.INT)), + schema -> epochDays, + value -> assertEquals(localDate, value) + ); + } + + // Test method for converting a single time field (number of milliseconds from midnight) + @Test + public void testTimeConversion() { + LocalTime localTime = LocalTime.of(10, 0); + long epochMicros = localTime.toNanoOfDay() / 1000; + int epochMillis = (int) (localTime.toNanoOfDay() / 1_000_000); + assertFieldRoundTrips("TimeMicros", "timeField", + () -> LogicalTypes.timeMicros().addToSchema(Schema.create(Schema.Type.LONG)), + schema -> epochMicros, + value -> assertEquals(localTime, value) + ); + + assertFieldRoundTrips("TimeMillis", "timeField2", + () -> LogicalTypes.timeMillis().addToSchema(Schema.create(Schema.Type.INT)), + schema -> epochMillis, + value -> assertEquals(localTime, value) + ); + } + + // Test method for converting a single timestamp field (number of milliseconds from epoch) + // timestamp: Stores microseconds from 1970-01-01 00:00:00.000000. [1] + // timestamptz: Stores microseconds from 1970-01-01 00:00:00.000000 UTC. [1] + @Test + public void testTimestampConversion() { + Instant instant = Instant.parse("2020-01-01T12:34:56.123456Z"); + long timestampMicros = instant.getEpochSecond() * 1_000_000 + instant.getNano() / 1_000; + long timestampMillis = instant.toEpochMilli(); + + Supplier timestampMicrosTzSchema = () -> { + Schema schema = LogicalTypes.timestampMicros().addToSchema(Schema.create(Schema.Type.LONG)); + schema.addProp("adjust-to-utc", true); + return schema; + }; + + Supplier timestampMicrosSchema = () -> { + Schema schema = LogicalTypes.timestampMicros().addToSchema(Schema.create(Schema.Type.LONG)); + schema.addProp("adjust-to-utc", false); + return schema; + }; + + Supplier timestampMillisTzSchema = () -> { + Schema schema = LogicalTypes.timestampMillis().addToSchema(Schema.create(Schema.Type.LONG)); + schema.addProp("adjust-to-utc", true); + return schema; + }; + + Supplier timestampMillisSchema = () -> { + Schema schema = LogicalTypes.timestampMillis().addToSchema(Schema.create(Schema.Type.LONG)); + schema.addProp("adjust-to-utc", false); + return schema; + }; + + OffsetDateTime expectedMicrosTz = DateTimeUtil.timestamptzFromMicros(timestampMicros); + LocalDateTime expectedMicros = DateTimeUtil.timestampFromMicros(timestampMicros); + OffsetDateTime expectedMillisTz = DateTimeUtil.timestamptzFromMicros(timestampMillis * 1000); + LocalDateTime expectedMillis = DateTimeUtil.timestampFromMicros(timestampMillis * 1000); + + assertFieldRoundTrips("TimestampMicrosTz", "timestampField1", + timestampMicrosTzSchema, + schema -> timestampMicros, + value -> assertEquals(expectedMicrosTz, value) + ); + + assertFieldRoundTrips("TimestampMicros", "timestampField2", + timestampMicrosSchema, + schema -> timestampMicros, + value -> assertEquals(expectedMicros, value) + ); + + assertFieldRoundTrips("TimestampMillisTz", "timestampField3", + timestampMillisTzSchema, + schema -> timestampMillis, + value -> assertEquals(expectedMillisTz, value) + ); + + assertFieldRoundTrips("TimestampMillis", "timestampField4", + timestampMillisSchema, + schema -> timestampMillis, + value -> assertEquals(expectedMillis, value) + ); + } + + @Test + public void testLocalTimestampConversion() { + LocalDateTime localDateTime = LocalDateTime.of(2023, 6, 1, 8, 15, 30, 123456000); + long micros = DateTimeUtil.microsFromTimestamp(localDateTime); + long millis = DateTimeUtil.microsToMillis(micros); + + // For millis precision, we need to truncate to milliseconds + LocalDateTime localDateTimeMillis = DateTimeUtil.timestampFromMicros(millis * 1000); + + Supplier localTimestampMillisSchema = () -> + LogicalTypes.localTimestampMillis().addToSchema(Schema.create(Schema.Type.LONG)); + Supplier localTimestampMicrosSchema = () -> + LogicalTypes.localTimestampMicros().addToSchema(Schema.create(Schema.Type.LONG)); + + assertFieldRoundTrips("LocalTimestampMillis", "localTsMillis", + localTimestampMillisSchema, + schema -> millis, + value -> { + assertEquals(millis, value); + assertEquals(localDateTimeMillis, DateTimeUtil.timestampFromMicros(((Long) value) * 1000)); + } + ); + + assertFieldRoundTrips("LocalTimestampMicros", "localTsMicros", + localTimestampMicrosSchema, + schema -> micros, + value -> { + assertEquals(micros, value); + assertEquals(localDateTime, DateTimeUtil.timestampFromMicros((Long) value)); + } + ); + } + + // Test method for converting a single binary field + @Test + public void testBinaryConversion() { + String randomAlphabetic = RandomStringUtils.randomAlphabetic(64); + assertFieldRoundTrips("Binary", "binaryField", + () -> Schema.create(Schema.Type.BYTES), + schema -> ByteBuffer.wrap(randomAlphabetic.getBytes(StandardCharsets.UTF_8)), + value -> { + ByteBuffer binaryField = (ByteBuffer) value; + assertEquals(randomAlphabetic, new String(binaryField.array(), StandardCharsets.UTF_8)); + } + ); + } + + // Test method for converting a single fixed field + @Test + public void testFixedConversion() { + assertFieldRoundTrips("Fixed", "fixedField", + () -> Schema.createFixed("FixedField", null, null, 3), + schema -> new GenericData.Fixed(schema, "bar".getBytes(StandardCharsets.UTF_8)), + value -> assertEquals("bar", new String((byte[]) value, StandardCharsets.UTF_8)) + ); + } + + // Test method for converting a single enum field + @Test + public void testEnumConversion() { + assertFieldRoundTrips("Enum", "enumField", + () -> Schema.createEnum("EnumField", null, null, Arrays.asList("A", "B", "C")), + schema -> new GenericData.EnumSymbol(schema, "B"), + value -> assertEquals("B", value.toString()) + ); + } + + // Test method for converting a single UUID field + @Test + public void testUUIDConversion() { + UUID uuid = UUID.randomUUID(); + assertFieldRoundTrips("UUID", "uuidField", + () -> LogicalTypes.uuid().addToSchema(Schema.create(Schema.Type.STRING)), + schema -> new Conversions.UUIDConversion().toCharSequence(uuid, schema, LogicalTypes.uuid()), + value -> assertEquals(uuid, UUIDUtil.convert((byte[]) value)) + ); + } + + // Test method for converting a single decimal field + @Test + public void testDecimalConversion() { + BigDecimal bigDecimal = BigDecimal.valueOf(1000.00).setScale(2); + assertFieldRoundTrips("Decimal", "decimalField", + () -> LogicalTypes.decimal(9, 2).addToSchema(Schema.create(Schema.Type.BYTES)), + schema -> { + LogicalTypes.Decimal decimalType = (LogicalTypes.Decimal) schema.getLogicalType(); + return new Conversions.DecimalConversion().toBytes(bigDecimal, schema, decimalType); + }, + value -> assertEquals(bigDecimal, value) + ); + } + + @Test + public void testStructFieldConversion() { + Schema structSchema = SchemaBuilder.record("NestedStruct") + .fields() + .name("field1").type().stringType().noDefault() + .name("field2").type().intType().noDefault() + .endRecord(); + + GenericRecord expected = new GenericData.Record(structSchema); + expected.put("field1", "nested_value"); + expected.put("field2", 99); + + assertFieldRoundTrips("StructField", "structField", + () -> structSchema, + schema -> cloneStruct(expected, schema), + value -> assertStructEquals(expected, (Record) value) + ); + } + + // Test method for converting a list field + @Test + public void testListConversion() { + List expected = Arrays.asList("a", "b", "c"); + assertFieldRoundTrips("List", "listField", + () -> Schema.createArray(Schema.create(Schema.Type.STRING)), + schema -> new ArrayList<>(expected), + value -> assertEquals(expected, normalizeValue(value)) + ); + } + + // Test method for converting a list of structs + @Test + public void testListStructConversion() { + Schema structSchema = SchemaBuilder.record("Struct") + .fields() + .name("field1").type().stringType().noDefault() + .name("field2").type().intType().noDefault() + .endRecord(); + + List expectedList = new ArrayList<>(); + + GenericRecord struct1 = new GenericData.Record(structSchema); + struct1.put("field1", "value1"); + struct1.put("field2", 1); + expectedList.add(struct1); + + GenericRecord struct2 = new GenericData.Record(structSchema); + struct2.put("field1", "value2"); + struct2.put("field2", 2); + expectedList.add(struct2); + + assertFieldRoundTrips("StructList", "listField", + () -> Schema.createArray(structSchema), + schema -> new ArrayList<>(expectedList), + value -> assertStructListEquals(expectedList, value) + ); + } + + // Test method for converting a list with nullable elements + @Test + public void testListWithNullableElementsConversion() { + assertFieldRoundTrips("ListNullableElements", "listField", + () -> Schema.createArray(Schema.createUnion(Arrays.asList( + Schema.create(Schema.Type.NULL), + Schema.create(Schema.Type.STRING) + ))), + schema -> { + @SuppressWarnings("unchecked") + GenericData.Array listValue = new GenericData.Array<>(3, schema); + listValue.add(new Utf8("a")); + listValue.add(null); + listValue.add(new Utf8("c")); + return listValue; + }, + value -> assertEquals(Arrays.asList("a", null, "c"), normalizeValue(value)) + ); + } + + @Test + public void testMapWithNonStringKeysConversion() { + Map expected = new LinkedHashMap<>(); + expected.put(1, "one"); + expected.put(2, "two"); + + Schema logicalMapSchema = createLogicalMapSchema("IntStringEntry", + Schema.create(Schema.Type.INT), Schema.create(Schema.Type.STRING)); + + assertFieldRoundTrips("IntKeyLogicalMap", "mapField", + () -> logicalMapSchema, + schema -> createLogicalMapArrayValue(schema, expected), + value -> { + Map actual = (Map) value; + Map normalized = new LinkedHashMap<>(); + actual.forEach((k, v) -> normalized.put((Integer) k, v == null ? null : v.toString())); + assertEquals(expected, normalized); + } + ); + } + + // Test method for converting a map with string values + @Test + public void testStringMapConversion() { + Map map = new HashMap<>(); + map.put("key1", "value1"); + map.put("key2", "value2"); + assertFieldRoundTrips("StringMap", "mapField", + () -> Schema.createMap(Schema.create(Schema.Type.STRING)), + schema -> new HashMap<>(map), + value -> assertEquals(map, normalizeValue(value)) + ); + } + + // Test method for converting a map with integer values + @Test + public void testIntMapConversion() { + Map map = new HashMap<>(); + map.put("key1", 1); + map.put("key2", 2); + assertFieldRoundTrips("IntMap", "mapField", + () -> Schema.createMap(Schema.create(Schema.Type.INT)), + schema -> new HashMap<>(map), + value -> assertEquals(map, normalizeValue(value)) + ); + } + + // Test method for converting a map with struct values + @Test + public void testStructMapConversion() { + Schema structSchema = SchemaBuilder.record("Struct") + .fields() + .name("field1").type().stringType().noDefault() + .name("field2").type().intType().noDefault() + .endRecord(); + + Map map = new HashMap<>(); + GenericRecord struct1 = new GenericData.Record(structSchema); + struct1.put("field1", "value1"); + struct1.put("field2", 1); + map.put("key1", struct1); + + GenericRecord struct2 = new GenericData.Record(structSchema); + struct2.put("field1", "value2"); + struct2.put("field2", 2); + map.put("key2", struct2); + + assertFieldRoundTrips("StructMap", "mapField", + () -> Schema.createMap(structSchema), + schema -> new HashMap<>(map), + value -> assertStructMapEquals(map, value) + ); + } + + // Test method for converting a map with nullable values + @Test + public void testMapWithNullableValuesConversion() { + Map expectedMap = new HashMap<>(); + expectedMap.put("key1", "value1"); + expectedMap.put("key2", null); + + assertFieldRoundTrips("NullableValueMap", "mapField", + () -> Schema.createMap(Schema.createUnion(Arrays.asList( + Schema.create(Schema.Type.NULL), + Schema.create(Schema.Type.STRING) + ))), + schema -> new HashMap<>(expectedMap), + value -> assertEquals(expectedMap, normalizeValue(value)) + ); + } + + + @Test + public void testBinaryFieldBackedByFixedConversion() { + Schema fixedSchema = Schema.createFixed("FixedBinary", null, null, 4); + Schema recordSchema = SchemaBuilder.builder() + .record("FixedBinaryRecord") + .namespace(TEST_NAMESPACE) + .fields() + .name("binaryField").type(fixedSchema).noDefault() + .endRecord(); + + Types.StructType structType = Types.StructType.of( + Types.NestedField.required(1, "binaryField", Types.BinaryType.get()) + ); + org.apache.iceberg.Schema icebergSchema = new org.apache.iceberg.Schema(structType.fields()); + + runRoundTrip(recordSchema, icebergSchema, + record -> record.put("binaryField", new GenericData.Fixed(fixedSchema, new byte[]{1, 2, 3, 4})), + icebergRecord -> { + ByteBuffer buffer = (ByteBuffer) icebergRecord.getField("binaryField"); + byte[] actual = new byte[buffer.remaining()]; + buffer.get(actual); + assertArrayEquals(new byte[]{1, 2, 3, 4}, actual); + } + ); + } + + // Test method for deeply nested struct (3+ levels) + @Test + public void testDeeplyNestedStructConversion() { + Schema innerMostStruct = SchemaBuilder.record("InnerMostStruct") + .namespace(TEST_NAMESPACE) + .fields() + .name("deepValue").type().intType().noDefault() + .endRecord(); + + Schema middleStruct = SchemaBuilder.record("MiddleStruct") + .namespace(TEST_NAMESPACE) + .fields() + .name("middleField").type().stringType().noDefault() + .name("innerMost").type(innerMostStruct).noDefault() + .endRecord(); + + Schema outerStruct = SchemaBuilder.record("OuterStruct") + .namespace(TEST_NAMESPACE) + .fields() + .name("outerField").type().stringType().noDefault() + .name("middle").type(middleStruct).noDefault() + .endRecord(); + + Schema recordSchema = SchemaBuilder.builder() + .record("DeeplyNestedRecord") + .namespace(TEST_NAMESPACE) + .fields() + .name("topLevel").type().stringType().noDefault() + .name("nested").type(outerStruct).noDefault() + .endRecord(); + + GenericRecord innerMostRecord = new GenericData.Record(innerMostStruct); + innerMostRecord.put("deepValue", 42); + + GenericRecord middleRecord = new GenericData.Record(middleStruct); + middleRecord.put("middleField", "middle"); + middleRecord.put("innerMost", innerMostRecord); + + GenericRecord outerRecord = new GenericData.Record(outerStruct); + outerRecord.put("outerField", "outer"); + outerRecord.put("middle", middleRecord); + + runRoundTrip(recordSchema, + record -> { + record.put("topLevel", "top"); + record.put("nested", outerRecord); + }, + icebergRecord -> { + assertEquals("top", icebergRecord.getField("topLevel").toString()); + Record nestedRecord = (Record) icebergRecord.getField("nested"); + assertNotNull(nestedRecord); + assertEquals("outer", nestedRecord.getField("outerField").toString()); + + Record middleResult = (Record) nestedRecord.getField("middle"); + assertNotNull(middleResult); + assertEquals("middle", middleResult.getField("middleField").toString()); + + Record innerMostResult = (Record) middleResult.getField("innerMost"); + assertNotNull(innerMostResult); + assertEquals(42, innerMostResult.getField("deepValue")); + } + ); + } + + // Test method for converting a record with default values + @Test + public void testDefaultFieldConversion() { + Schema recordSchema = SchemaBuilder.builder() + .record("DefaultValueRecord") + .namespace(TEST_NAMESPACE) + .fields() + .name("defaultStringField").type().stringType().stringDefault("default_string") + .name("defaultIntField").type().intType().intDefault(42) + .name("defaultBoolField").type().booleanType().booleanDefault(true) + .endRecord(); + + // Test with default values + runRoundTrip(recordSchema, + record -> { + Schema.Field defaultStringField = recordSchema.getField("defaultStringField"); + Schema.Field defaultIntField = recordSchema.getField("defaultIntField"); + Schema.Field defaultBoolField = recordSchema.getField("defaultBoolField"); + record.put("defaultStringField", defaultStringField.defaultVal()); + record.put("defaultIntField", defaultIntField.defaultVal()); + record.put("defaultBoolField", defaultBoolField.defaultVal()); + }, + icebergRecord -> { + assertEquals("default_string", icebergRecord.getField("defaultStringField").toString()); + assertEquals(42, icebergRecord.getField("defaultIntField")); + assertEquals(true, icebergRecord.getField("defaultBoolField")); + } + ); + + // Test with non-default values + runRoundTrip(recordSchema, + record -> { + record.put("defaultStringField", "custom_value"); + record.put("defaultIntField", 100); + record.put("defaultBoolField", false); + }, + icebergRecord -> { + assertEquals("custom_value", icebergRecord.getField("defaultStringField").toString()); + assertEquals(100, icebergRecord.getField("defaultIntField")); + assertEquals(false, icebergRecord.getField("defaultBoolField")); + } + ); + } + + // Test that non-optional unions with multiple non-NULL types throw UnsupportedOperationException + @Test + public void testNonOptionalUnionThrowsException() { + // Test case 1: {null, string, int} at record level + Schema unionSchema1 = Schema.createUnion(Arrays.asList( + Schema.create(Schema.Type.NULL), + Schema.create(Schema.Type.STRING), + Schema.create(Schema.Type.INT) + )); + + try { + RecordBinder binder = new RecordBinder(AvroSchemaUtil.toIceberg(unionSchema1), unionSchema1); + org.junit.jupiter.api.Assertions.fail("Expected UnsupportedOperationException for non-optional union {null, string, int}"); + } catch (UnsupportedOperationException e) { + assertEquals(true, e.getMessage().contains("Non-optional UNION with multiple non-NULL types is not supported")); + assertEquals(true, e.getMessage().contains("Found 2 non-NULL types")); + } + + // Test case 2: {null, struct1, struct2} at record level + Schema struct1Schema = SchemaBuilder.record("Struct1") + .namespace(TEST_NAMESPACE) + .fields() + .name("field1").type().stringType().noDefault() + .endRecord(); + + Schema struct2Schema = SchemaBuilder.record("Struct2") + .namespace(TEST_NAMESPACE) + .fields() + .name("field2").type().intType().noDefault() + .endRecord(); + + Schema unionSchema2 = Schema.createUnion(Arrays.asList( + Schema.create(Schema.Type.NULL), + struct1Schema, + struct2Schema + )); + + try { + RecordBinder binder = new RecordBinder(AvroSchemaUtil.toIceberg(unionSchema2), unionSchema2); + org.junit.jupiter.api.Assertions.fail("Expected UnsupportedOperationException for non-optional union {null, struct1, struct2}"); + } catch (UnsupportedOperationException e) { + assertEquals(true, e.getMessage().contains("Non-optional UNION with multiple non-NULL types is not supported")); + assertEquals(true, e.getMessage().contains("Found 2 non-NULL types")); + } + + // Test case 3: Union in field with multiple non-NULL types + Schema unionFieldSchema = Schema.createUnion(Arrays.asList( + Schema.create(Schema.Type.NULL), + Schema.create(Schema.Type.STRING), + Schema.create(Schema.Type.INT) + )); + + Schema recordSchema = SchemaBuilder.builder() + .record("RecordWithUnionField") + .namespace(TEST_NAMESPACE) + .fields() + .name("id").type().intType().noDefault() + .name("unionField").type(unionFieldSchema).withDefault(null) + .endRecord(); + + try { + RecordBinder binder = new RecordBinder(AvroSchemaUtil.toIceberg(recordSchema), recordSchema); + org.junit.jupiter.api.Assertions.fail("Expected UnsupportedOperationException for field with non-optional union"); + } catch (UnsupportedOperationException e) { + assertEquals(true, e.getMessage().contains("Non-optional UNION with multiple non-NULL types is not supported")); + assertEquals(true, e.getMessage().contains("Found 2 non-NULL types")); + } + } + + + private void testSendRecord(org.apache.iceberg.Schema schema, org.apache.iceberg.data.Record record) { + String tableName = "test_" + tableCounter++; + table = catalog.createTable(TableIdentifier.of(Namespace.of("default"), tableName), schema); + writer = createTableWriter(table); + try { + writer.write(record); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + public static TaskWriter createTableWriter(Table table) { + FileAppenderFactory appenderFactory = new GenericAppenderFactory( + table.schema(), + table.spec(), + null, null, null) + .setAll(new HashMap<>(table.properties())) + .set(PARQUET_ROW_GROUP_SIZE_BYTES, "1"); + + OutputFileFactory fileFactory = + OutputFileFactory.builderFor(table, 1, System.currentTimeMillis()) + .defaultSpec(table.spec()) + .operationId(UUID.randomUUID().toString()) + .format(FileFormat.PARQUET) + .build(); + + return new UnpartitionedWriter<>( + table.spec(), + FileFormat.PARQUET, + appenderFactory, + fileFactory, + table.io(), + WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT + ); + } + + private static GenericRecord serializeAndDeserialize(GenericRecord record, Schema schema) { + try { + // Serialize the avro record to a byte array + ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + DatumWriter datumWriter = new SpecificDatumWriter<>(schema); + Encoder encoder = EncoderFactory.get().binaryEncoder(outputStream, null); + datumWriter.write(record, encoder); + encoder.flush(); + outputStream.close(); + + byte[] serializedBytes = outputStream.toByteArray(); + + // Deserialize the byte array back to an avro record + DatumReader datumReader = new SpecificDatumReader<>(schema); + ByteArrayInputStream inputStream = new ByteArrayInputStream(serializedBytes); + Decoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null); + return datumReader.read(null, decoder); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + + private static Schema createOptionalSchema(Schema schema) { + if (schema.getType() == Schema.Type.UNION) { + boolean hasNull = schema.getTypes().stream() + .anyMatch(type -> type.getType() == Schema.Type.NULL); + if (hasNull) { + return schema; + } + List updatedTypes = new ArrayList<>(); + updatedTypes.add(Schema.create(Schema.Type.NULL)); + updatedTypes.addAll(schema.getTypes()); + return Schema.createUnion(updatedTypes); + } + return Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), schema)); + } + + private static Schema ensureNonNullBranch(Schema schema) { + if (schema.getType() != Schema.Type.UNION) { + return schema; + } + return schema.getTypes().stream() + .filter(type -> type.getType() != Schema.Type.NULL) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("Union schema lacks non-null branch: " + schema)); + } + + private void runRoundTrip(Schema recordSchema, Consumer avroPopulator, Consumer assertions) { + runRoundTrip(recordSchema, AvroSchemaUtil.toIceberg(recordSchema), avroPopulator, assertions); + } + + private void runRoundTrip(Schema recordSchema, + org.apache.iceberg.Schema icebergSchema, + Consumer avroPopulator, + Consumer assertions) { + GenericRecord avroRecord = new GenericData.Record(recordSchema); + avroPopulator.accept(avroRecord); + GenericRecord roundTripRecord = serializeAndDeserialize(avroRecord, recordSchema); + + Record icebergRecord = new RecordBinder(icebergSchema, recordSchema).bind(roundTripRecord); + + assertions.accept(icebergRecord); + testSendRecord(icebergSchema, icebergRecord); + } + + // Helper method to test round-trip conversion for a single field + private void assertFieldRoundTrips(String recordPrefix, + String fieldName, + Supplier fieldSchemaSupplier, + Function avroValueSupplier, + Consumer valueAssertion) { + Schema baseFieldSchema = fieldSchemaSupplier.get(); + Schema baseRecordSchema = SchemaBuilder.builder() + .record(recordPrefix + "Base") + .namespace(TEST_NAMESPACE) + .fields() + .name(fieldName).type(baseFieldSchema).noDefault() + .endRecord(); + + // Direct field + runRoundTrip(baseRecordSchema, + record -> record.put(fieldName, avroValueSupplier.apply(baseFieldSchema)), + icebergRecord -> valueAssertion.accept(icebergRecord.getField(fieldName)) + ); + + Schema optionalFieldSchema = createOptionalSchema(fieldSchemaSupplier.get()); + Schema unionRecordSchema = SchemaBuilder.builder() + .record(recordPrefix + "Union") + .namespace(TEST_NAMESPACE) + .fields() + .name(fieldName).type(optionalFieldSchema).withDefault(null) + .endRecord(); + Schema nonNullBranch = ensureNonNullBranch(optionalFieldSchema); + + // Optional field with non-null value + runRoundTrip(unionRecordSchema, + record -> record.put(fieldName, avroValueSupplier.apply(nonNullBranch)), + icebergRecord -> valueAssertion.accept(icebergRecord.getField(fieldName)) + ); + + // Optional field with null value + runRoundTrip(unionRecordSchema, + record -> record.put(fieldName, null), + icebergRecord -> assertNull(icebergRecord.getField(fieldName)) + ); + } + + + private static Map toStringKeyMap(Object value) { + if (value == null) { + return null; + } + Map map = (Map) value; + Map result = new HashMap<>(map.size()); + for (Map.Entry entry : map.entrySet()) { + String key = entry.getKey() == null ? null : entry.getKey().toString(); + result.put(key, normalizeValue(entry.getValue())); + } + return result; + } + + private static GenericRecord cloneStruct(GenericRecord source, Schema schema) { + GenericRecord target = new GenericData.Record(schema); + for (Schema.Field field : schema.getFields()) { + target.put(field.name(), source.get(field.name())); + } + return target; + } + + private static Schema createLogicalMapSchema(String entryName, Schema keySchema, Schema valueSchema) { + Schema.Field keyField = new Schema.Field("key", keySchema, null, null); + Schema.Field valueField = new Schema.Field("value", valueSchema, null, null); + Schema entrySchema = Schema.createRecord(entryName, null, null, false); + entrySchema.setFields(Arrays.asList(keyField, valueField)); + Schema arraySchema = Schema.createArray(entrySchema); + return CodecSetup.getLogicalMap().addToSchema(arraySchema); + } + + private static GenericData.Array createLogicalMapArrayValue(Schema schema, Map values) { + Schema nonNullSchema = ensureNonNullBranch(schema); + if (nonNullSchema.getType() != Schema.Type.ARRAY) { + throw new IllegalArgumentException("Expected array schema for logical map but got: " + nonNullSchema); + } + Schema entrySchema = nonNullSchema.getElementType(); + Schema.Field keyField = entrySchema.getField("key"); + Schema.Field valueField = entrySchema.getField("value"); + GenericData.Array entries = new GenericData.Array<>(values.size(), nonNullSchema); + for (Map.Entry entry : values.entrySet()) { + GenericRecord kv = new GenericData.Record(entrySchema); + kv.put(keyField.name(), toAvroValue(entry.getKey(), keyField.schema())); + kv.put(valueField.name(), toAvroValue(entry.getValue(), valueField.schema())); + entries.add(kv); + } + return entries; + } + + private static Object toAvroValue(Object value, Schema schema) { + if (value == null) { + return null; + } + Schema actualSchema = ensureNonNullBranch(schema); + switch (actualSchema.getType()) { + case STRING: + return value instanceof CharSequence ? value : new Utf8(value.toString()); + case INT: + case LONG: + case FLOAT: + case DOUBLE: + case BOOLEAN: + return value; + case RECORD: + return value; + default: + return value; + } + } + + private static List toRecordList(Object value) { + if (value == null) { + return null; + } + List list = (List) value; + List normalized = new ArrayList<>(list.size()); + for (Object element : list) { + normalized.add((Record) element); + } + return normalized; + } + + private static Map toRecordMap(Object value) { + if (value == null) { + return null; + } + Map map = (Map) value; + Map normalized = new HashMap<>(map.size()); + for (Map.Entry entry : map.entrySet()) { + String key = entry.getKey() == null ? null : entry.getKey().toString(); + normalized.put(key, (Record) entry.getValue()); + } + return normalized; + } + + private static void assertStructListEquals(List expectedList, Object actualValue) { + List actualList = toRecordList(actualValue); + assertNotNull(actualList, "Actual list is null"); + assertEquals(expectedList.size(), actualList.size()); + for (int i = 0; i < expectedList.size(); i++) { + assertStructEquals(expectedList.get(i), actualList.get(i)); + } + } + + private static void assertStructMapEquals(Map expectedMap, Object actualValue) { + Map actualMap = toRecordMap(actualValue); + assertNotNull(actualMap, "Actual map is null"); + assertEquals(expectedMap.keySet(), actualMap.keySet()); + for (Map.Entry entry : expectedMap.entrySet()) { + assertStructEquals(entry.getValue(), actualMap.get(entry.getKey())); + } + } + + private static void assertStructEquals(GenericRecord expected, Record actual) { + assertNotNull(actual, "Actual struct record is null"); + for (Schema.Field field : expected.getSchema().getFields()) { + Object expectedValue = normalizeValue(expected.get(field.name())); + Object actualValue = normalizeValue(actual.getField(field.name())); + assertEquals(expectedValue, actualValue, "Mismatch on field " + field.name()); + } + } + + private static Object normalizeValue(Object value) { + if (value == null) { + return null; + } + if (value instanceof CharSequence) { + return value.toString(); + } + if (value instanceof List) { + List list = (List) value; + List normalized = new ArrayList<>(list.size()); + for (Object element : list) { + normalized.add(normalizeValue(element)); + } + return normalized; + } + if (value instanceof Map) { + return toStringKeyMap(value); + } + return value; + } + +} diff --git a/core/src/test/java/kafka/automq/table/coordinator/CheckpointTest.java b/core/src/test/java/kafka/automq/table/coordinator/CheckpointTest.java new file mode 100644 index 0000000000..4f7f333018 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/coordinator/CheckpointTest.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.coordinator; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +@Tag("S3Unit") +public class CheckpointTest { + + @Test + public void testCodec() { + Checkpoint checkpoint = new Checkpoint(Status.REQUEST_COMMIT, UUID.randomUUID(), 10L, new long[] {2, 3, 3}, UUID.randomUUID(), 1000L, new long[] {1, 2, 3}); + Checkpoint rst = Checkpoint.decode(ByteBuffer.wrap(checkpoint.encode())); + assertEquals(checkpoint.status(), rst.status()); + assertEquals(checkpoint.commitId(), rst.commitId()); + assertEquals(checkpoint.taskOffset(), rst.taskOffset()); + assertArrayEquals(checkpoint.nextOffsets(), checkpoint.nextOffsets()); + assertEquals(checkpoint.lastCommitId(), rst.lastCommitId()); + assertEquals(checkpoint.lastCommitTimestamp(), rst.lastCommitTimestamp()); + assertArrayEquals(checkpoint.preCommitOffsets(), rst.preCommitOffsets()); + } + +} diff --git a/core/src/test/java/kafka/automq/table/coordinator/TableCoordinatorTest.java b/core/src/test/java/kafka/automq/table/coordinator/TableCoordinatorTest.java new file mode 100644 index 0000000000..b5749701a2 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/coordinator/TableCoordinatorTest.java @@ -0,0 +1,588 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.coordinator; + +import kafka.automq.table.Channel; +import kafka.automq.table.events.CommitRequest; +import kafka.automq.table.events.CommitResponse; +import kafka.automq.table.events.Envelope; +import kafka.automq.table.events.Errors; +import kafka.automq.table.events.Event; +import kafka.automq.table.events.EventType; +import kafka.automq.table.events.PartitionMetric; +import kafka.automq.table.events.TopicMetric; +import kafka.automq.table.events.WorkerOffset; +import kafka.automq.table.utils.PartitionUtil; +import kafka.automq.table.utils.TableIdentifierUtil; +import kafka.log.streamaspect.MetaKeyValue; +import kafka.log.streamaspect.MetaStream; +import kafka.server.MetadataCache; + +import org.apache.kafka.common.config.TopicConfig; +import org.apache.kafka.common.message.UpdateMetadataRequestData; +import org.apache.kafka.storage.internals.log.LogConfig; + +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.ExpireSnapshots; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.Transaction; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.ArrayDeque; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; + +import scala.Option; +import scala.Some; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class TableCoordinatorTest { + + private static final String TOPIC = "test-topic"; + + @Mock + Channel channel; + @Mock + MetaStream metaStream; + @Mock + MetadataCache metadataCache; + + private InMemoryCatalog catalog; + private Table table; + private TableCoordinator coordinator; + private Supplier configSupplier; + private TableCoordinator.CommitStatusMachine machine; + private FakeSubChannel subChannel; + + @BeforeEach + void setUp() { + LogConfig logConfig = new FakeLogConfig(1000L, "db", ""); + configSupplier = () -> logConfig; + + // metadata stubs + UpdateMetadataRequestData.UpdateMetadataPartitionState state = new UpdateMetadataRequestData.UpdateMetadataPartitionState(); + state.setLeaderEpoch(1); + doReturn((Option) Some.apply(2)).when(metadataCache).numPartitions(TOPIC); + when(metadataCache.getPartitionInfo(eq(TOPIC), anyInt())).thenReturn(Option.apply(state)); + + // in-memory iceberg catalog & table + Schema schema = new Schema(Types.NestedField.required(1, "id", Types.IntegerType.get())); + PartitionSpec spec = PartitionUtil.buildPartitionSpec(List.of(), schema); + TableIdentifier identifier = TableIdentifierUtil.of("db", TOPIC); + catalog = new InMemoryCatalog(); + catalog.initialize("test", Map.of()); + catalog.createNamespace(identifier.namespace()); + table = catalog.createTable(identifier, schema, spec); + + // meta stream stub + when(metaStream.append(any(MetaKeyValue.class))).thenReturn(CompletableFuture.completedFuture(new DummyAppendResult())); + + // channel stub + subChannel = new FakeSubChannel(); + when(channel.subscribeData(eq(TOPIC), anyLong())).thenAnswer(invocation -> subChannel); + + coordinator = new TableCoordinator(catalog, TOPIC, metaStream, channel, new ImmediateEventLoop(), metadataCache, configSupplier); + machine = coordinator.new CommitStatusMachine(); + } + + @Test + void nextRoundCommitSendsCommitRequestAndCheckpoint() throws Exception { + machine.nextRoundCommit(); + + ArgumentCaptor eventCaptor = ArgumentCaptor.forClass(Event.class); + verify(channel).send(eq(TOPIC), eventCaptor.capture()); + Event event = eventCaptor.getValue(); + assertEquals(EventType.COMMIT_REQUEST, event.type()); + CommitRequest payload = event.payload(); + assertNotNull(payload.commitId()); + assertEquals(2, payload.offsets().size()); + + verify(metaStream).append(any(MetaKeyValue.class)); + assertEquals(Status.REQUEST_COMMIT, machine.status); + } + + @Test + void commitResponseMovesToCommittedAndWritesIcebergSnapshot() throws Exception { + machine.nextRoundCommit(); + + UUID commitId = machine.processing.commitId; + Types.StructType partitionType = table.spec().partitionType(); + List nextOffsets = List.of(new WorkerOffset(0, 1, 5L), new WorkerOffset(1, 1, 6L)); + DataFile dataFile = DataFiles.builder(table.spec()) + .withPath("file:///tmp/commit.parquet") + .withFileSizeInBytes(10) + .withRecordCount(1) + .build(); + CommitResponse response = new CommitResponse(partitionType, Errors.NONE, commitId, TOPIC, nextOffsets, + List.of(dataFile), List.of(), new TopicMetric(10), + List.of(new PartitionMetric(0, 100L), new PartitionMetric(1, 200L))); + subChannel.offer(new Envelope(0, 3L, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, response))); + + machine.tryMoveToCommittedStatus(); + table.refresh(); + + assertEquals(Status.COMMITTED, machine.status); + assertArrayEquals(new long[]{5L, 6L}, machine.last.nextOffsets); + Snapshot snapshot = table.currentSnapshot(); + assertNotNull(snapshot); + assertEquals(commitId.toString(), snapshot.summary().get("automq.commit.id")); + assertEquals("100", snapshot.summary().get("automq.watermark")); + } + + @Test + void moreDataResponseEnablesFastNextCommit() throws Exception { + machine.nextRoundCommit(); + UUID commitId = machine.processing.commitId; + Types.StructType partitionType = table.spec().partitionType(); + List nextOffsets = List.of(new WorkerOffset(0, 1, 2L), new WorkerOffset(1, 1, 3L)); + CommitResponse response = new CommitResponse(partitionType, Errors.MORE_DATA, commitId, TOPIC, nextOffsets, List.of(), List.of(), TopicMetric.NOOP, List.of()); + subChannel.offer(new Envelope(0, 1L, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, response))); + + machine.tryMoveToCommittedStatus(); + + assertTrue(getPrivateBoolean(machine, "fastNextCommit")); + } + + @Test + void watermarkHelperWorks() { + assertEquals(-1L, TableCoordinator.watermark(new long[]{-1L, -1L})); + assertEquals(100L, TableCoordinator.watermark(new long[]{100L, 200L})); + assertEquals(50L, TableCoordinator.watermark(new long[]{-1L, 50L})); + } + + @Test + void commitTimesOutButStillAdvances() throws Exception { + machine.nextRoundCommit(); + setPrivateLong(machine, "requestCommitTimestamp", System.currentTimeMillis() - 60_000); + + machine.tryMoveToCommittedStatus(); + + assertEquals(Status.COMMITTED, machine.status); + assertTrue(getPrivateBoolean(machine, "fastNextCommit")); + } + + @Test + void checkpointRecoveryFromRequestCommitInitializesState() throws Exception { + UUID commitId = UUID.randomUUID(); + long[] next = new long[]{1L, 2L}; + Checkpoint cp = new Checkpoint(Status.REQUEST_COMMIT, commitId, 10L, next, UUID.randomUUID(), 0L, new long[0]); + TableCoordinator.CommitStatusMachine recovered = coordinator.new CommitStatusMachine(cp); + + assertEquals(Status.REQUEST_COMMIT, recovered.status); + assertArrayEquals(next, recovered.processing.nextOffsets); + assertEquals(2, getPartitionWatermarks(recovered).length); + } + + @Test + void checkpointPreCommitSkipsAlreadyCommittedSnapshot() { + UUID commitId = UUID.randomUUID(); + table.newAppend() + .set("automq.commit.id", commitId.toString()) + .set("automq.watermark", "123") + .commit(); + + long[] next = new long[]{3L, 4L}; + Checkpoint cp = new Checkpoint(Status.PRE_COMMIT, commitId, 5L, next, UUID.randomUUID(), 0L, new long[]{3L, 4L}); + TableCoordinator.CommitStatusMachine recovered = coordinator.new CommitStatusMachine(cp); + + assertEquals(Status.COMMITTED, recovered.status); + assertArrayEquals(next, recovered.last.nextOffsets); + assertEquals(commitId, recovered.last.commitId); + } + + @Test + void partitionNumIncreaseExpandsArrays() throws Exception { + machine.nextRoundCommit(); // init with 2 partitions + doReturn((Option) Some.apply(4)).when(metadataCache).numPartitions(TOPIC); + + machine.nextRoundCommit(); + + assertEquals(4, machine.processing.nextOffsets.length); + assertEquals(4, getPartitionWatermarks(machine).length); + } + + @Test + void partitionByEvolutionTriggersEvolve() throws Exception { + setPrivateField(coordinator, "table", table); + setLogConfigField("tableTopicPartitionBy", "id"); + + Method evolve = machine.getClass().getDeclaredMethod("tryEvolvePartition"); + evolve.setAccessible(true); + boolean evolved = (boolean) evolve.invoke(machine); + + assertTrue(evolved); + } + + @Test + void expireSnapshotsHonorsDefaultRetention() throws Exception { + SpyHolder spyHolder = spyTableForExpireVerification(table); + setPrivateField(coordinator, "table", spyHolder.tableSpy); + + machine.nextRoundCommit(); + UUID commitId = machine.processing.commitId; + + DataFile dataFile = DataFiles.builder(table.spec()) + .withPath("file:///tmp/commit.parquet") + .withFileSizeInBytes(10) + .withRecordCount(1) + .build(); + + CommitResponse response = createCommitResponse(commitId, List.of(dataFile)); + subChannel.offer(new Envelope(0, 1L, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, response))); + + machine.tryMoveToCommittedStatus(); + + assertEquals(Status.COMMITTED, machine.status); + verifyExpireSnapshotsCalledWith(spyHolder.capturedExpireSnapshots, 1, 1); + } + + @Test + void expireSnapshotsUsesConfiguredValues() throws Exception { + SpyHolder spyHolder = spyTableForExpireVerification(table); + setPrivateField(coordinator, "table", spyHolder.tableSpy); + setCustomExpireConfig(5, 3, true); + + machine.nextRoundCommit(); + UUID commitId = machine.processing.commitId; + + DataFile dataFile = DataFiles.builder(table.spec()) + .withPath("file:///tmp/commit.parquet") + .withFileSizeInBytes(10) + .withRecordCount(1) + .build(); + + CommitResponse response = createCommitResponse(commitId, List.of(dataFile)); + subChannel.offer(new Envelope(0, 1L, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, response))); + + machine.tryMoveToCommittedStatus(); + + assertEquals(Status.COMMITTED, machine.status); + verifyExpireSnapshotsCalledWith(spyHolder.capturedExpireSnapshots, 3, 5); + } + + @Test + void expireSnapshotsDisabledSkipsCall() throws Exception { + SpyHolder spyHolder = spyTableForExpireVerification(table); + setPrivateField(coordinator, "table", spyHolder.tableSpy); + setCustomExpireConfig(0, 0, false); + + machine.nextRoundCommit(); + UUID commitId = machine.processing.commitId; + + CommitResponse response = createCommitResponse(commitId, List.of()); + subChannel.offer(new Envelope(0, 0L, new Event(System.currentTimeMillis(), EventType.COMMIT_RESPONSE, response))); + + machine.tryMoveToCommittedStatus(); + + assertEquals(Status.COMMITTED, machine.status); + verify(spyHolder.tableSpy, Mockito.never()).newTransaction(); + } + + // --- test helpers --- + private CommitResponse createCommitResponse(UUID commitId, List dataFiles) { + Types.StructType partitionType = table.spec().partitionType(); + List nextOffsets = List.of(new WorkerOffset(0, 1, 5L), new WorkerOffset(1, 1, 6L)); + TopicMetric topicMetric = dataFiles.isEmpty() ? TopicMetric.NOOP : new TopicMetric(1); + List partitionMetrics = List.of(new PartitionMetric(0, 10L), new PartitionMetric(1, 20L)); + + return new CommitResponse(partitionType, Errors.NONE, commitId, TOPIC, nextOffsets, + dataFiles, List.of(), topicMetric, partitionMetrics); + } + + private void setCustomExpireConfig(int olderThanHours, int retainLast, boolean enabled) throws Exception { + Map props = new HashMap<>(); + props.put(TopicConfig.TABLE_TOPIC_ENABLE_CONFIG, true); + props.put(TopicConfig.TABLE_TOPIC_COMMIT_INTERVAL_CONFIG, 1000L); + props.put(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG, "db"); + props.put(TopicConfig.SEGMENT_BYTES_CONFIG, 1073741824); + props.put(TopicConfig.RETENTION_MS_CONFIG, 86400000L); + props.put(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_CONFIG, enabled); + if (enabled) { + props.put(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_CONFIG, olderThanHours); + props.put(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_CONFIG, retainLast); + } + Supplier custom = () -> new LogConfig(props); + setPrivateField(coordinator, "config", custom); + } + + private void verifyExpireSnapshotsCalledWith(ExpireSnapshots expireSnapshots, int retainLast, int olderThanHours) { + assertNotNull(expireSnapshots, "ExpireSnapshots should have been captured"); + + verify(expireSnapshots).retainLast(retainLast); + + ArgumentCaptor olderThanCaptor = ArgumentCaptor.forClass(Long.class); + verify(expireSnapshots).expireOlderThan(olderThanCaptor.capture()); + long expectedOlderThan = System.currentTimeMillis() - TimeUnit.HOURS.toMillis(olderThanHours); + long actualOlderThan = olderThanCaptor.getValue(); + assertTrue(Math.abs(actualOlderThan - expectedOlderThan) < TimeUnit.SECONDS.toMillis(5), + String.format("Expected olderThan within 5s of %d hours ago, but was %d ms off", + olderThanHours, Math.abs(actualOlderThan - expectedOlderThan))); + + verify(expireSnapshots).executeDeleteWith(any()); + verify(expireSnapshots).commit(); + } + + private static boolean getPrivateBoolean(Object target, String name) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + return field.getBoolean(target); + } + + private static class FakeSubChannel implements Channel.SubChannel { + private final ArrayDeque queue = new ArrayDeque<>(); + + void offer(Envelope envelope) { + queue.offer(envelope); + } + + @Override + public Envelope poll() { + return queue.poll(); + } + + @Override + public void close() { + } + } + + private static class ImmediateEventLoop extends com.automq.stream.utils.threads.EventLoop { + ImmediateEventLoop() { + super("immediate-loop"); + } + + @Override + public void execute(Runnable command) { + command.run(); + } + + @Override + public java.util.concurrent.CompletableFuture submit(Runnable task) { + task.run(); + return CompletableFuture.completedFuture(null); + } + } + + private static class FakeLogConfig extends LogConfig { + FakeLogConfig(long commitInterval, String namespace, String partitionBy) { + super(buildProps(commitInterval, namespace, partitionBy)); + } + + private static Map buildProps(long commitInterval, String namespace, String partitionBy) { + Map props = new HashMap<>(); + props.put(TopicConfig.TABLE_TOPIC_ENABLE_CONFIG, true); + props.put(TopicConfig.TABLE_TOPIC_COMMIT_INTERVAL_CONFIG, commitInterval); + props.put(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG, namespace); + props.put(TopicConfig.TABLE_TOPIC_PARTITION_BY_CONFIG, partitionBy); + // supply required basics to satisfy defaults + props.put(TopicConfig.SEGMENT_BYTES_CONFIG, 1073741824); + props.put(TopicConfig.RETENTION_MS_CONFIG, 86400000L); + return props; + } + } + + private static class DummyAppendResult implements com.automq.stream.api.AppendResult { + @Override + public long baseOffset() { + return 0; + } + } + + private static void setPrivateLong(Object target, String name, long value) throws Exception { + Field f = target.getClass().getDeclaredField(name); + f.setAccessible(true); + f.setLong(target, value); + } + + private static void setPrivateField(Object target, String name, Object value) throws Exception { + Field f = target.getClass().getDeclaredField(name); + f.setAccessible(true); + f.set(target, value); + } + + private void setLogConfigField(String name, Object value) throws Exception { + Field f = LogConfig.class.getDeclaredField(name); + f.setAccessible(true); + f.set(configSupplier.get(), value); + } + + private static long[] getPartitionWatermarks(TableCoordinator.CommitStatusMachine machine) throws Exception { + Field f = machine.getClass().getDeclaredField("partitionWatermarks"); + f.setAccessible(true); + return (long[]) f.get(machine); + } + + // --- Spy infrastructure for testing ExpireSnapshots --- + + private static class SpyHolder { + final Table tableSpy; + volatile ExpireSnapshots capturedExpireSnapshots; + + SpyHolder(Table tableSpy) { + this.tableSpy = tableSpy; + } + } + + private static SpyHolder spyTableForExpireVerification(Table delegate) { + Table tableSpy = Mockito.spy(delegate); + SpyHolder holder = new SpyHolder(tableSpy); + + Mockito.doAnswer(invocation -> { + Transaction realTxn = (Transaction) invocation.callRealMethod(); + return new TransactionWrapper(realTxn, holder); + }).when(tableSpy).newTransaction(); + + return holder; + } + + /** + * Transparent wrapper for Transaction that only intercepts expireSnapshots() + * to create a spy for verification purposes. + */ + private static class TransactionWrapper implements Transaction { + private final Transaction delegate; + private final SpyHolder holder; + + TransactionWrapper(Transaction delegate, SpyHolder holder) { + this.delegate = delegate; + this.holder = holder; + } + + @Override + public ExpireSnapshots expireSnapshots() { + ExpireSnapshots realExpire = delegate.expireSnapshots(); + ExpireSnapshots expireSpy = Mockito.spy(realExpire); + holder.capturedExpireSnapshots = expireSpy; + return expireSpy; + } + + // All other methods delegate transparently + @Override + public org.apache.iceberg.AppendFiles newAppend() { + return delegate.newAppend(); + } + + @Override + public org.apache.iceberg.AppendFiles newFastAppend() { + return delegate.newFastAppend(); + } + + @Override + public org.apache.iceberg.RewriteFiles newRewrite() { + return delegate.newRewrite(); + } + + @Override + public org.apache.iceberg.RewriteManifests rewriteManifests() { + return delegate.rewriteManifests(); + } + + @Override + public org.apache.iceberg.OverwriteFiles newOverwrite() { + return delegate.newOverwrite(); + } + + @Override + public org.apache.iceberg.RowDelta newRowDelta() { + return delegate.newRowDelta(); + } + + @Override + public org.apache.iceberg.ReplacePartitions newReplacePartitions() { + return delegate.newReplacePartitions(); + } + + @Override + public org.apache.iceberg.DeleteFiles newDelete() { + return delegate.newDelete(); + } + + @Override + public org.apache.iceberg.UpdateProperties updateProperties() { + return delegate.updateProperties(); + } + + @Override + public org.apache.iceberg.UpdateSchema updateSchema() { + return delegate.updateSchema(); + } + + @Override + public org.apache.iceberg.UpdatePartitionSpec updateSpec() { + return delegate.updateSpec(); + } + + @Override + public org.apache.iceberg.UpdateLocation updateLocation() { + return delegate.updateLocation(); + } + + @Override + public org.apache.iceberg.ReplaceSortOrder replaceSortOrder() { + return delegate.replaceSortOrder(); + } + + @Override + public void commitTransaction() { + delegate.commitTransaction(); + } + + @Override + public org.apache.iceberg.Table table() { + return delegate.table(); + } + } +} diff --git a/core/src/test/java/kafka/automq/table/deserializer/proto/parse/ProtobufSchemaWellKnownTypesTest.java b/core/src/test/java/kafka/automq/table/deserializer/proto/parse/ProtobufSchemaWellKnownTypesTest.java new file mode 100644 index 0000000000..b7fa82faa4 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/deserializer/proto/parse/ProtobufSchemaWellKnownTypesTest.java @@ -0,0 +1,402 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.deserializer.proto.parse; + +import kafka.automq.table.deserializer.proto.schema.ProtobufSchema; + +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; + +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Collections; +import java.util.Locale; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests that verify parsing schema files with Google Well-Known Types and other complex protobuf features. + * + * Note: These tests are currently disabled because they require access to the Well-Known Type definition files + * which are not bundled with the test resources. In a real environment, these would be available + * through the protobuf library. To make these tests pass, you would need to: + * 1. Include the well-known type .proto files in your test resources + * 2. Make them available to the ProtobufSchemaParser when parsing + * + * Note: The current implementation does not support the `extend` feature. + */ +@Tag("S3Unit") +class ProtobufSchemaWellKnownTypesTest { + + private static final String TEST_RESOURCES_DIR = "src/test/resources/proto"; + + @Test + @DisplayName("Should parse schema with google.protobuf.Timestamp") + void shouldParseSchemaWithTimestamp() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_timestamp.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + Descriptors.Descriptor messageDescriptor = fileDescriptor.findMessageTypeByName("EventWithTimestamp"); + assertNotNull(messageDescriptor); + + // Verify timestamp fields + Descriptors.FieldDescriptor createdAtField = messageDescriptor.findFieldByName("created_at"); + Descriptors.FieldDescriptor updatedAtField = messageDescriptor.findFieldByName("updated_at"); + + assertNotNull(createdAtField); + assertNotNull(updatedAtField); + + assertEquals("google.protobuf.Timestamp", createdAtField.getMessageType().getFullName()); + assertEquals("google.protobuf.Timestamp", updatedAtField.getMessageType().getFullName()); + + // Verify field types + assertEquals(Descriptors.FieldDescriptor.Type.MESSAGE, createdAtField.getType()); + assertEquals(Descriptors.FieldDescriptor.Type.MESSAGE, updatedAtField.getType()); + } + + @Test + @DisplayName("Should parse schema with multiple well-known types") + void shouldParseSchemaWithWellKnownTypes() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_well_known_types.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + Descriptors.Descriptor messageDescriptor = fileDescriptor.findMessageTypeByName("ComplexTypeMessage"); + assertNotNull(messageDescriptor); + + // Verify imports are preserved + DescriptorProtos.FileDescriptorProto fileDescriptorProto = fileDescriptor.toProto(); + assertTrue(fileDescriptorProto.getDependencyList().contains("google/protobuf/timestamp.proto")); + assertTrue(fileDescriptorProto.getDependencyList().contains("google/protobuf/duration.proto")); + assertTrue(fileDescriptorProto.getDependencyList().contains("google/protobuf/any.proto")); + assertTrue(fileDescriptorProto.getDependencyList().contains("google/protobuf/struct.proto")); + + // Verify well-known type fields + Map fieldTypeMap = messageDescriptor.getFields().stream() + .filter(f -> f.getType() == Descriptors.FieldDescriptor.Type.MESSAGE) + .collect(Collectors.toMap( + Descriptors.FieldDescriptor::getName, + f -> f.getMessageType().getFullName() + )); + + assertEquals("google.protobuf.Timestamp", fieldTypeMap.get("created_at")); + assertEquals("google.protobuf.Duration", fieldTypeMap.get("elapsed_time")); + assertEquals("google.protobuf.Any", fieldTypeMap.get("details")); + assertEquals("google.protobuf.Struct", fieldTypeMap.get("attributes")); + assertEquals("google.protobuf.StringValue", fieldTypeMap.get("optional_name")); + assertEquals("google.protobuf.BoolValue", fieldTypeMap.get("is_active")); + assertEquals("google.protobuf.Int64Value", fieldTypeMap.get("big_count")); + assertEquals("google.protobuf.DoubleValue", fieldTypeMap.get("score")); + assertEquals("google.protobuf.FieldMask", fieldTypeMap.get("update_mask")); + assertEquals("google.protobuf.Empty", fieldTypeMap.get("nothing")); + } + + @Test + @DisplayName("Should parse schema with map fields") + void shouldParseSchemaWithMapFields() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_map.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + Descriptors.Descriptor messageDescriptor = fileDescriptor.findMessageTypeByName("MapMessage"); + assertNotNull(messageDescriptor); + + // Verify map fields are correctly identified as repeated messages with key/value fields + Descriptors.FieldDescriptor stringToStringField = messageDescriptor.findFieldByName("string_to_string"); + Descriptors.FieldDescriptor intToStringField = messageDescriptor.findFieldByName("int_to_string"); + Descriptors.FieldDescriptor stringToNestedField = messageDescriptor.findFieldByName("string_to_nested"); + + assertNotNull(stringToStringField); + assertNotNull(intToStringField); + assertNotNull(stringToNestedField); + + // Maps are represented as repeated messages with a special entry type + assertTrue(stringToStringField.isMapField()); + assertTrue(intToStringField.isMapField()); + assertTrue(stringToNestedField.isMapField()); + + // Verify the entry types + Descriptors.Descriptor stringToStringEntryType = stringToStringField.getMessageType(); + assertEquals("string_to_stringentry", stringToStringEntryType.getName().toLowerCase(Locale.ENGLISH)); + assertEquals(Descriptors.FieldDescriptor.Type.STRING, + stringToStringEntryType.findFieldByName("key").getType()); + assertEquals(Descriptors.FieldDescriptor.Type.STRING, + stringToStringEntryType.findFieldByName("value").getType()); + + // Verify map with nested message value + Descriptors.Descriptor stringToNestedEntryType = stringToNestedField.getMessageType(); + Descriptors.FieldDescriptor valueField = stringToNestedEntryType.findFieldByName("value"); + assertEquals(Descriptors.FieldDescriptor.Type.MESSAGE, valueField.getType()); + assertEquals("NestedValue", valueField.getMessageType().getName()); + } + + @Test + @DisplayName("Should parse schema with repeated fields") + void shouldParseSchemaWithRepeatedFields() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_repeated_fields.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + Descriptors.Descriptor messageDescriptor = fileDescriptor.findMessageTypeByName("RepeatedFieldsMessage"); + assertNotNull(messageDescriptor); + + // Verify repeated fields + Descriptors.FieldDescriptor tagsField = messageDescriptor.findFieldByName("tags"); + Descriptors.FieldDescriptor valuesField = messageDescriptor.findFieldByName("values"); + Descriptors.FieldDescriptor itemsField = messageDescriptor.findFieldByName("items"); + Descriptors.FieldDescriptor eventTimesField = messageDescriptor.findFieldByName("event_times"); + + assertTrue(tagsField.isRepeated()); + assertTrue(valuesField.isRepeated()); + assertTrue(itemsField.isRepeated()); + assertTrue(eventTimesField.isRepeated()); + + // Verify repeated of complex message types + assertEquals(Descriptors.FieldDescriptor.Type.MESSAGE, itemsField.getType()); + assertEquals("Item", itemsField.getMessageType().getName()); + + // Verify repeated of well-known types + assertEquals(Descriptors.FieldDescriptor.Type.MESSAGE, eventTimesField.getType()); + assertEquals("google.protobuf.Timestamp", eventTimesField.getMessageType().getFullName()); + } + + @Test + @DisplayName("Should parse schema with reserved fields") + void shouldParseSchemaWithReservedFields() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_reserved.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + Descriptors.Descriptor messageDescriptor = fileDescriptor.findMessageTypeByName("MessageWithReserved"); + assertNotNull(messageDescriptor); + + // Get the proto representation to verify reserved fields + DescriptorProtos.DescriptorProto descriptorProto = messageDescriptor.toProto(); + + // Check reserved numbers + assertTrue(descriptorProto.getReservedRangeList().stream() + .anyMatch(range -> range.getStart() == 2 && range.getEnd() == 2)); // Single number 2 + assertTrue(descriptorProto.getReservedRangeList().stream() + .anyMatch(range -> range.getStart() == 15 && range.getEnd() == 15)); // Single number 15 + assertTrue(descriptorProto.getReservedRangeList().stream() + .anyMatch(range -> range.getStart() == 9 && range.getEnd() == 11)); // Range 9-11 + assertTrue(descriptorProto.getReservedRangeList().stream() + .anyMatch(range -> range.getStart() == 40 && range.getEnd() == 45)); // Range 40-45 + + // Check reserved names + assertTrue(descriptorProto.getReservedNameList().contains("foo")); + assertTrue(descriptorProto.getReservedNameList().contains("bar")); + assertTrue(descriptorProto.getReservedNameList().contains("baz")); + } + + @Test + @DisplayName("Should parse schema with service definitions") + @Disabled + void shouldParseSchemaWithServiceDefinitions() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_service.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + DescriptorProtos.FileDescriptorProto fileDescriptorProto = fileDescriptor.toProto(); + + // Verify service is present + assertEquals(1, fileDescriptorProto.getServiceCount()); + assertEquals("UserService", fileDescriptorProto.getService(0).getName()); + + // Verify the service methods + DescriptorProtos.ServiceDescriptorProto serviceProto = fileDescriptorProto.getService(0); + assertEquals(8, serviceProto.getMethodCount()); + + // Verify specific methods and their streaming properties + Map methods = serviceProto.getMethodList().stream() + .collect(Collectors.toMap( + DescriptorProtos.MethodDescriptorProto::getName, + method -> method + )); + + // Regular unary method + assertTrue(methods.containsKey("GetUser")); + assertFalse(methods.get("GetUser").getClientStreaming()); + assertFalse(methods.get("GetUser").getServerStreaming()); + + // Server streaming method + assertTrue(methods.containsKey("StreamUserUpdates")); + assertFalse(methods.get("StreamUserUpdates").getClientStreaming()); + assertTrue(methods.get("StreamUserUpdates").getServerStreaming()); + + // Client streaming method + assertTrue(methods.containsKey("UploadUserData")); + assertTrue(methods.get("UploadUserData").getClientStreaming()); + assertFalse(methods.get("UploadUserData").getServerStreaming()); + + // Bidirectional streaming method + assertTrue(methods.containsKey("ProcessUserBatch")); + assertTrue(methods.get("ProcessUserBatch").getClientStreaming()); + assertTrue(methods.get("ProcessUserBatch").getServerStreaming()); + } + + @Test + @DisplayName("Should parse schema with extensions") + @Disabled + void shouldParseSchemaWithExtensions() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_extensions.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + + // Verify the base message + Descriptors.Descriptor baseMessageDescriptor = fileDescriptor.findMessageTypeByName("ExtendingMessage"); + assertNotNull(baseMessageDescriptor); + + // Verify extension ranges in base message + DescriptorProtos.DescriptorProto baseProto = baseMessageDescriptor.toProto(); + assertFalse(baseProto.getExtensionRangeList().isEmpty()); + + // Verify extension range values + assertTrue(baseProto.getExtensionRangeList().stream() + .anyMatch(range -> range.getStart() == 100 && range.getEnd() == 199)); + } + + @Test + @DisplayName("Should parse schema with complex nested types") + void shouldParseSchemaWithNestedTypes() throws Exception { + // Given + Path protoPath = Paths.get(TEST_RESOURCES_DIR, "with_nested_types.proto"); + String protoContent = Files.readString(protoPath); + + // When + ProtobufSchema schema = ProtobufSchemaParser.parseSchema(protoContent, Collections.emptyMap()); + + // Then + Descriptors.FileDescriptor fileDescriptor = schema.getFileDescriptor(); + + // Verify the main message type + Descriptors.Descriptor mainMessageDescriptor = fileDescriptor.findMessageTypeByName("ComplexNestedMessage"); + assertNotNull(mainMessageDescriptor); + + // Verify nested enum + Descriptors.EnumDescriptor statusEnum = mainMessageDescriptor.findEnumTypeByName("Status"); + assertNotNull(statusEnum); + assertEquals(4, statusEnum.getValues().size()); + assertEquals("UNKNOWN", statusEnum.getValues().get(0).getName()); + assertEquals("ACTIVE", statusEnum.getValues().get(1).getName()); + + // Verify nested Address message + Descriptors.Descriptor addressDescriptor = mainMessageDescriptor.findNestedTypeByName("Address"); + assertNotNull(addressDescriptor); + assertEquals(4, addressDescriptor.getFields().size()); + + // Verify deeply nested GeoLocation message + Descriptors.Descriptor geoLocationDescriptor = addressDescriptor.findNestedTypeByName("GeoLocation"); + assertNotNull(geoLocationDescriptor); + assertEquals(3, geoLocationDescriptor.getFields().size()); + + // Verify deepest nested Accuracy message + Descriptors.Descriptor accuracyDescriptor = geoLocationDescriptor.findNestedTypeByName("Accuracy"); + assertNotNull(accuracyDescriptor); + assertEquals(2, accuracyDescriptor.getFields().size()); + + // Verify recursive TreeNode message + Descriptors.Descriptor treeNodeDescriptor = mainMessageDescriptor.findNestedTypeByName("TreeNode"); + assertNotNull(treeNodeDescriptor); + + // Verify recursive field in TreeNode + Descriptors.FieldDescriptor childrenField = treeNodeDescriptor.findFieldByName("children"); + assertNotNull(childrenField); + assertTrue(childrenField.isRepeated()); + assertEquals(Descriptors.FieldDescriptor.Type.MESSAGE, childrenField.getType()); + assertEquals(treeNodeDescriptor, childrenField.getMessageType()); + + // Verify ContactInfo with oneof + Descriptors.Descriptor contactInfoDescriptor = mainMessageDescriptor.findNestedTypeByName("ContactInfo"); + assertNotNull(contactInfoDescriptor); + + // Verify oneof fields + Descriptors.OneofDescriptor oneofDescriptor = contactInfoDescriptor.getOneofs().get(0); + assertEquals("contact", oneofDescriptor.getName()); + assertEquals(3, oneofDescriptor.getFields().size()); + assertTrue(oneofDescriptor.getFields().stream() + .anyMatch(field -> field.getName().equals("email"))); + assertTrue(oneofDescriptor.getFields().stream() + .anyMatch(field -> field.getName().equals("phone"))); + assertTrue(oneofDescriptor.getFields().stream() + .anyMatch(field -> field.getName().equals("physical_address"))); + + // Verify map fields + Descriptors.FieldDescriptor labeledAddressesField = mainMessageDescriptor.findFieldByName("labeled_addresses"); + assertNotNull(labeledAddressesField); + assertTrue(labeledAddressesField.isMapField()); + assertEquals(Descriptors.FieldDescriptor.Type.STRING, labeledAddressesField.getMessageType().findFieldByName("key").getType()); + assertEquals(addressDescriptor, labeledAddressesField.getMessageType().findFieldByName("value").getMessageType()); + + // Verify repeated fields + Descriptors.FieldDescriptor secondaryAddressesField = mainMessageDescriptor.findFieldByName("secondary_addresses"); + assertNotNull(secondaryAddressesField); + assertTrue(secondaryAddressesField.isRepeated()); + assertEquals(addressDescriptor, secondaryAddressesField.getMessageType()); + } +} diff --git a/core/src/test/java/kafka/automq/table/events/EventTest.java b/core/src/test/java/kafka/automq/table/events/EventTest.java new file mode 100644 index 0000000000..720e8e500c --- /dev/null +++ b/core/src/test/java/kafka/automq/table/events/EventTest.java @@ -0,0 +1,94 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.List; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@Tag("S3Unit") +public class EventTest { + + @Test + public void testCommitRequestCodec() throws IOException { + Event event = new Event(1, EventType.COMMIT_REQUEST, + new CommitRequest(UUID.randomUUID(), "topic-xx", List.of(new WorkerOffset(1, 2, 3))) + ); + Event rst = AvroCodec.decode(AvroCodec.encode(event)); + + assertEquals(event.timestamp(), rst.timestamp()); + assertEquals(event.type(), rst.type()); + CommitRequest req1 = event.payload(); + CommitRequest req2 = rst.payload(); + assertEquals(req1.commitId(), req2.commitId()); + assertEquals(req1.topic(), req2.topic()); + assertEquals(req1.offsets().size(), req2.offsets().size()); + assertTrue(workOffsetEqual(req1.offsets().get(0), req2.offsets().get(0))); + } + + @Test + public void testCommitResponseCodec() throws IOException { + Event event = new Event(2, EventType.COMMIT_RESPONSE, + new CommitResponse( + Types.StructType.of(), + 233, + UUID.randomUUID(), + "topic", + List.of(new WorkerOffset(1, 2, 3)), + List.of(EventTestUtil.createDataFile()), List.of(EventTestUtil.createDeleteFile(), + EventTestUtil.createDeleteFile()), + new TopicMetric(233), + List.of(new PartitionMetric(1, 200)) + ) + ); + Event rst = AvroCodec.decode(AvroCodec.encode(event)); + assertEquals(event.timestamp(), rst.timestamp()); + assertEquals(event.type(), rst.type()); + CommitResponse resp1 = event.payload(); + CommitResponse resp2 = rst.payload(); + assertEquals(resp1.code(), resp2.code()); + assertEquals(resp1.commitId(), resp2.commitId()); + assertEquals(resp1.topic(), resp2.topic()); + assertEquals(resp1.nextOffsets().size(), resp2.nextOffsets().size()); + assertTrue(workOffsetEqual(resp1.nextOffsets().get(0), resp2.nextOffsets().get(0))); + + assertEquals(1, resp2.dataFiles().size()); + assertEquals(resp1.dataFiles().get(0).path(), resp2.dataFiles().get(0).path()); + + assertEquals(2, resp2.deleteFiles().size()); + assertEquals(resp1.deleteFiles().get(0).path(), resp2.deleteFiles().get(0).path()); + assertEquals(resp1.deleteFiles().get(1).path(), resp2.deleteFiles().get(1).path()); + + assertEquals(resp1.topicMetric(), resp2.topicMetric()); + assertEquals(resp1.partitionMetrics(), resp2.partitionMetrics()); + } + + private boolean workOffsetEqual(WorkerOffset o1, WorkerOffset o2) { + return o1.partition() == o2.partition() && o1.epoch() == o2.epoch() && o1.offset() == o2.offset(); + } + +} diff --git a/core/src/test/java/kafka/automq/table/events/EventTestUtil.java b/core/src/test/java/kafka/automq/table/events/EventTestUtil.java new file mode 100644 index 0000000000..b97c875234 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/events/EventTestUtil.java @@ -0,0 +1,54 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.events; + +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.PartitionSpec; +import org.junit.jupiter.api.Tag; + +import java.util.Map; +import java.util.UUID; + +@Tag("S3Unit") +public class EventTestUtil { + public static DataFile createDataFile() { + return DataFiles.builder(PartitionSpec.unpartitioned()) + .withPath(UUID.randomUUID() + ".parquet") + .withFormat(FileFormat.PARQUET) + .withFileSizeInBytes(100L) + .withRecordCount(5) + .withMetrics(new Metrics(1L, Map.of(1, 2L, 3, 4L), null, null, null)) + .build(); + } + + public static DeleteFile createDeleteFile() { + return FileMetadata.deleteFileBuilder(PartitionSpec.unpartitioned()) + .ofEqualityDeletes(1) + .withPath(UUID.randomUUID() + ".parquet") + .withFileSizeInBytes(10) + .withRecordCount(1) + .build(); + } +} \ No newline at end of file diff --git a/core/src/test/java/kafka/automq/table/process/DefaultRecordProcessorTest.java b/core/src/test/java/kafka/automq/table/process/DefaultRecordProcessorTest.java new file mode 100644 index 0000000000..c7d3683cde --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/DefaultRecordProcessorTest.java @@ -0,0 +1,677 @@ + +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.convert.AvroRegistryConverter; +import kafka.automq.table.process.convert.RawConverter; +import kafka.automq.table.process.convert.StringConverter; +import kafka.automq.table.process.exception.ConverterException; +import kafka.automq.table.process.exception.InvalidDataException; +import kafka.automq.table.process.exception.SchemaRegistrySystemException; +import kafka.automq.table.process.exception.TransformException; +import kafka.automq.table.process.transform.FlattenTransform; + +import org.apache.kafka.common.cache.Cache; +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.header.internals.RecordHeader; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.TimestampType; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; + +import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; +import io.confluent.kafka.serializers.KafkaAvroSerializer; + +import static java.net.HttpURLConnection.HTTP_NOT_FOUND; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@Tag("S3Unit") +public class DefaultRecordProcessorTest { + + private static final String TEST_TOPIC = "test-topic"; + private static final int TEST_PARTITION = 0; + private static final long TEST_OFFSET = 123L; + private static final long TEST_TIMESTAMP = System.currentTimeMillis(); + + private SchemaRegistryClient schemaRegistryClient; + private KafkaAvroSerializer avroSerializer; + + private static final Schema USER_SCHEMA_V1 = SchemaBuilder.record("User") + .namespace("kafka.automq.table.process") + .fields() + .name("name").type().stringType().noDefault() + .endRecord(); + + private static final Schema USER_SCHEMA_V2 = SchemaBuilder.record("User") + .namespace("kafka.automq.table.process") + .fields() + .name("name").type().stringType().noDefault() + .name("age").type().intType().intDefault(0) + .endRecord(); + + @BeforeEach + void setUp() { + schemaRegistryClient = new MockSchemaRegistryClient(); + avroSerializer = new KafkaAvroSerializer(schemaRegistryClient); + avroSerializer.configure(Map.of("schema.registry.url", "http://mock:8081"), false); + } + + private Record createKafkaRecord(byte[] key, byte[] value, Header[] headers) { + return new SimpleRecord(TEST_OFFSET, TEST_TIMESTAMP, key, value, headers); + } + + private Record createAvroRecord(String topic, Object avroRecord, String key) { + byte[] value = avroSerializer.serialize(topic, avroRecord); + byte[] keyBytes = key.getBytes(); + return createKafkaRecord(keyBytes, value, new Header[0]); + } + + @Test + void testProcessWithSchemaUpdateShouldChangeSchemaIdentity() throws Exception { + // Arrange + String valueSubject = TEST_TOPIC + "-value"; + + // Register V1 schema + schemaRegistryClient.register(valueSubject, USER_SCHEMA_V1); + + // Create processor + Converter keyConverter = new StringConverter(); + Converter valueConverter = new AvroRegistryConverter(schemaRegistryClient, "http://mock:8081", false); + DefaultRecordProcessor recordProcessor = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter); + + // Create first record with V1 + GenericRecord userRecordV1A = new GenericRecordBuilder(USER_SCHEMA_V1) + .set("name", "test-user-A") + .build(); + Record kafkaRecordV1A = createAvroRecord(TEST_TOPIC, userRecordV1A, "key1A"); + + // Act 1 + ProcessingResult resultV1A = recordProcessor.process(TEST_PARTITION, kafkaRecordV1A); + + // Assert 1 + assertTrue(resultV1A.isSuccess()); + assertNotNull(resultV1A.getFinalSchemaIdentity()); + String identityV1A = resultV1A.getFinalSchemaIdentity(); + + // Arrange 2: Create second record with the same V1 schema + GenericRecord userRecordV1B = new GenericRecordBuilder(USER_SCHEMA_V1) + .set("name", "test-user-B") + .build(); + Record kafkaRecordV1B = createAvroRecord(TEST_TOPIC, userRecordV1B, "key1B"); + + // Act 2 + ProcessingResult resultV1B = recordProcessor.process(TEST_PARTITION, kafkaRecordV1B); + + // Assert 2 + assertTrue(resultV1B.isSuccess()); + String identityV1B = resultV1B.getFinalSchemaIdentity(); + assertEquals(identityV1A, identityV1B, "Schema identity should be the same for the same schema version"); + + // Arrange 3: Update schema to V2 + schemaRegistryClient.register(valueSubject, USER_SCHEMA_V2); + + // Create record with V2 + GenericRecord userRecordV2 = new GenericRecordBuilder(USER_SCHEMA_V2) + .set("name", "test-user-2") + .set("age", 30) + .build(); + Record kafkaRecordV2 = createAvroRecord(TEST_TOPIC, userRecordV2, "key2"); + + // Act 3 + ProcessingResult resultV2 = recordProcessor.process(TEST_PARTITION, kafkaRecordV2); + + // Assert 3 + assertTrue(resultV2.isSuccess()); + assertNotNull(resultV2.getFinalSchemaIdentity()); + String identityV2 = resultV2.getFinalSchemaIdentity(); + + // Final assertion + assertNotEquals(identityV1A, identityV2, "Schema identity should change after schema evolution"); + + GenericRecord finalRecordV2 = resultV2.getFinalRecord(); + GenericRecord valueRecordV2 = (GenericRecord) finalRecordV2.get(RecordAssembler.KAFKA_VALUE_FIELD); + assertEquals("test-user-2", valueRecordV2.get("name").toString()); + assertEquals(30, valueRecordV2.get("age")); + } + + @Test + void testBasicRawProcessing() { + // Arrange + Converter rawConverter = new RawConverter(); + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, rawConverter, rawConverter); + byte[] key = "test-key".getBytes(); + byte[] value = "test-value".getBytes(); + Record kafkaRecord = createKafkaRecord(key, value, new Header[0]); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + assertNull(result.getError()); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + assertEquals(ByteBuffer.wrap(key), ByteBuffer.wrap((byte[]) finalRecord.get(RecordAssembler.KAFKA_KEY_FIELD))); + assertEquals(ByteBuffer.wrap(value), ByteBuffer.wrap((byte[]) finalRecord.get(RecordAssembler.KAFKA_VALUE_FIELD))); + } + + @Test + void testConverterErrorHandling() { + // Arrange + Converter errorConverter = (topic, buffer) -> { + throw new ConverterException("Test conversion error"); + }; + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new StringConverter(), errorConverter); + Record kafkaRecord = createKafkaRecord("key".getBytes(), "value".getBytes(), new Header[0]); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertFalse(result.isSuccess()); + assertNotNull(result.getError()); + assertEquals(DataError.ErrorType.CONVERT_ERROR, result.getError().getType()); + assertTrue(result.getError().getMessage().contains("Test conversion error")); + } + + @Test + void testWithFlattenTransform() { + // Arrange + Converter keyConverter = new StringConverter(); + + Schema innerSchema = SchemaBuilder.record("Inner").fields().name("data").type().stringType().noDefault().endRecord(); + Converter valueConverter = (topic, buffer) -> { + GenericRecord innerRecord = new GenericRecordBuilder(innerSchema) + .set("data", "some-data") + .build(); + return new ConversionResult(innerRecord, "id1"); + }; + + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter, List.of(new FlattenTransform())); + Record kafkaRecord = createKafkaRecord("key".getBytes(), "value".getBytes(), new Header[0]); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + + // After flatten, the "data" field from innerRecord should be at the top level. + assertEquals("some-data", finalRecord.get("data").toString()); + assertEquals("key", finalRecord.get(RecordAssembler.KAFKA_KEY_FIELD)); + assertFalse(finalRecord.hasField("value"), "The original 'value' wrapper field should be gone after flatten."); + } + + @Test + void testHeaderProcessing() { + // Arrange + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), new RawConverter()); + Header[] headers = {new RecordHeader("h1", "v1".getBytes())}; + Record kafkaRecord = createKafkaRecord("key".getBytes(), "value".getBytes(), headers); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + @SuppressWarnings("unchecked") + Map headerMap = (Map) finalRecord.get(RecordAssembler.KAFKA_HEADER_FIELD); + assertEquals(1, headerMap.size()); + assertEquals(ByteBuffer.wrap("v1".getBytes()), headerMap.get("h1")); + } + + @Test + void testProcessHeadersWithMultipleEntriesIncludingNullValue() { + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), new RawConverter()); + Header[] headers = { + new RecordHeader("h1", "v1".getBytes(StandardCharsets.UTF_8)), + new RecordHeader("h2", null), + new RecordHeader("h3", "v3".getBytes(StandardCharsets.UTF_8)) + }; + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), headers); + + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + @SuppressWarnings("unchecked") + Map headerMap = (Map) finalRecord.get(RecordAssembler.KAFKA_HEADER_FIELD); + assertEquals(3, headerMap.size()); + assertEquals(ByteBuffer.wrap("v1".getBytes(StandardCharsets.UTF_8)), headerMap.get("h1")); + assertNull(headerMap.get("h2")); + assertEquals(ByteBuffer.wrap("v3".getBytes(StandardCharsets.UTF_8)), headerMap.get("h3")); + } + + @Test + void testProcessHeadersReuseEmptyResultInstance() { + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), new RawConverter()); + + Record recordWithoutHeaders = createKafkaRecord("key1".getBytes(StandardCharsets.UTF_8), "value1".getBytes(StandardCharsets.UTF_8), null); + ProcessingResult firstResult = processor.process(TEST_PARTITION, recordWithoutHeaders); + assertTrue(firstResult.isSuccess()); + @SuppressWarnings("unchecked") + Map firstHeaders = (Map) firstResult.getFinalRecord().get(RecordAssembler.KAFKA_HEADER_FIELD); + assertTrue(firstHeaders.isEmpty()); + + Record recordWithEmptyHeaders = createKafkaRecord("key2".getBytes(StandardCharsets.UTF_8), "value2".getBytes(StandardCharsets.UTF_8), new Header[0]); + ProcessingResult secondResult = processor.process(TEST_PARTITION, recordWithEmptyHeaders); + assertTrue(secondResult.isSuccess()); + @SuppressWarnings("unchecked") + Map secondHeaders = (Map) secondResult.getFinalRecord().get(RecordAssembler.KAFKA_HEADER_FIELD); + assertTrue(secondHeaders.isEmpty()); + assertSame(firstHeaders, secondHeaders); + } + + @Test + void testMetadataFieldsPopulatedOnSuccess() { + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), new RawConverter()); + int partition = 7; + long offset = 456L; + long timestamp = 1_234_567_890L; + Record kafkaRecord = new SimpleRecord(offset, timestamp, + "k".getBytes(StandardCharsets.UTF_8), "v".getBytes(StandardCharsets.UTF_8), new Header[0]); + + ProcessingResult result = processor.process(partition, kafkaRecord); + + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + GenericRecord metadata = (GenericRecord) finalRecord.get(RecordAssembler.KAFKA_METADATA_FIELD); + assertNotNull(metadata); + assertEquals(partition, ((Integer) metadata.get(RecordAssembler.METADATA_PARTITION_FIELD)).intValue()); + assertEquals(offset, ((Long) metadata.get(RecordAssembler.METADATA_OFFSET_FIELD)).longValue()); + assertEquals(timestamp, ((Long) metadata.get(RecordAssembler.METADATA_TIMESTAMP_FIELD)).longValue()); + } + + @Test + void testWrapValueSchemaCacheReusedBetweenCalls() { + Schema valueSchema = SchemaBuilder.record("CachedValue") + .namespace("kafka.automq.table.process.test") + .fields() + .name("field").type().stringType().noDefault() + .endRecord(); + AtomicBoolean alternateIdentity = new AtomicBoolean(false); + + Converter keyConverter = new StringConverter(); + Converter valueConverter = (topic, buffer) -> { + GenericRecord record = new GenericRecordBuilder(valueSchema) + .set("field", alternateIdentity.get() ? "value-b" : "value-a") + .build(); + String identity = alternateIdentity.get() ? "value-schema-b" : "value-schema-a"; + return new ConversionResult(record, identity); + }; + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter); + + Record kafkaRecord1 = createKafkaRecord("key1".getBytes(StandardCharsets.UTF_8), "value1".getBytes(StandardCharsets.UTF_8), new Header[0]); + ProcessingResult result1 = processor.process(TEST_PARTITION, kafkaRecord1); + assertTrue(result1.isSuccess()); + Cache cache = extractValueWrapperSchemaCache(processor); + assertEquals(1L, cache.size()); + + Record kafkaRecord2 = createKafkaRecord("key2".getBytes(StandardCharsets.UTF_8), "value2".getBytes(StandardCharsets.UTF_8), new Header[0]); + ProcessingResult result2 = processor.process(TEST_PARTITION, kafkaRecord2); + assertTrue(result2.isSuccess()); + assertEquals(1L, cache.size()); + + alternateIdentity.set(true); + Record kafkaRecord3 = createKafkaRecord("key3".getBytes(StandardCharsets.UTF_8), "value3".getBytes(StandardCharsets.UTF_8), new Header[0]); + ProcessingResult result3 = processor.process(TEST_PARTITION, kafkaRecord3); + assertTrue(result3.isSuccess()); + assertEquals(2L, cache.size()); + } + + @Test + void testCompositeSchemaIdentityReflectsTransformChain() { + Schema valueSchema = SchemaBuilder.record("IdentityRecord") + .namespace("kafka.automq.table.process.test") + .fields() + .name("field").type().stringType().noDefault() + .endRecord(); + Converter keyConverter = new StringConverter(); + Converter valueConverter = (topic, buffer) -> { + GenericRecord record = new GenericRecordBuilder(valueSchema) + .set("field", "payload") + .build(); + return new ConversionResult(record, "identity-value"); + }; + + DefaultRecordProcessor ordered = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter, + List.of(new NamedPassthroughTransform("A"), new NamedPassthroughTransform("B"))); + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), new Header[0]); + ProcessingResult orderedResult = ordered.process(TEST_PARTITION, kafkaRecord); + assertTrue(orderedResult.isSuccess()); + String orderedIdentity = orderedResult.getFinalSchemaIdentity(); + assertTrue(orderedIdentity.endsWith("|t:A,B")); + + DefaultRecordProcessor reversed = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter, + List.of(new NamedPassthroughTransform("B"), new NamedPassthroughTransform("A"))); + ProcessingResult reversedResult = reversed.process(TEST_PARTITION, kafkaRecord); + assertTrue(reversedResult.isSuccess()); + assertNotEquals(orderedIdentity, reversedResult.getFinalSchemaIdentity()); + } + + @Test + void testTransformReturningNullProducesError() { + Schema valueSchema = SchemaBuilder.record("NullRecord") + .namespace("kafka.automq.table.process.test") + .fields() + .name("field").type().stringType().noDefault() + .endRecord(); + Converter keyConverter = new StringConverter(); + Converter valueConverter = (topic, buffer) -> { + GenericRecord record = new GenericRecordBuilder(valueSchema) + .set("field", "payload") + .build(); + return new ConversionResult(record, "null-transform-identity"); + }; + + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter, List.of(new NullingTransform())); + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), new Header[0]); + + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + assertFalse(result.isSuccess()); + assertEquals(DataError.ErrorType.TRANSFORMATION_ERROR, result.getError().getType()); + assertTrue(result.getError().getMessage().contains("NullingTransform")); + } + + @Test + void testTransformThrowsInvalidDataException() { + Schema valueSchema = SchemaBuilder.record("InvalidRecord") + .namespace("kafka.automq.table.process.test") + .fields() + .name("field").type().stringType().noDefault() + .endRecord(); + Converter keyConverter = new StringConverter(); + Converter valueConverter = (topic, buffer) -> { + GenericRecord record = new GenericRecordBuilder(valueSchema) + .set("field", "payload") + .build(); + return new ConversionResult(record, "invalid-transform-identity"); + }; + + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter, List.of(new InvalidDataThrowingTransform())); + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), new Header[0]); + + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + assertFalse(result.isSuccess()); + assertEquals(DataError.ErrorType.DATA_ERROR, result.getError().getType()); + assertTrue(result.getError().getMessage().contains("Invalid data")); + } + + @Test + void testTransformThrowsTransformException() { + Schema valueSchema = SchemaBuilder.record("TransformRecord") + .namespace("kafka.automq.table.process.test") + .fields() + .name("field").type().stringType().noDefault() + .endRecord(); + Converter keyConverter = new StringConverter(); + Converter valueConverter = (topic, buffer) -> { + GenericRecord record = new GenericRecordBuilder(valueSchema) + .set("field", "payload") + .build(); + return new ConversionResult(record, "transform-exception-identity"); + }; + + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, keyConverter, valueConverter, List.of(new ThrowingTransform())); + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), new Header[0]); + + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + assertFalse(result.isSuccess()); + assertEquals(DataError.ErrorType.TRANSFORMATION_ERROR, result.getError().getType()); + assertTrue(result.getError().getMessage().contains("transform failure")); + } + + @Test + void testConverterRestClientNotFoundReturnsDataError() { + Converter restNotFoundConverter = (topic, buffer) -> { + RestClientException restException = new RestClientException("missing", HTTP_NOT_FOUND, 40403); + throw new RuntimeException("wrapper", restException); + }; + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), restNotFoundConverter); + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), null); + + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + assertFalse(result.isSuccess()); + assertEquals(DataError.ErrorType.CONVERT_ERROR, result.getError().getType()); + String errorMessage = result.getError().getMessage(); + assertNotNull(errorMessage); + assertTrue(errorMessage.contains("Schema or subject not found for record"), () -> "actual message: " + errorMessage); + assertTrue(errorMessage.contains("topic=" + TEST_TOPIC), () -> "actual message: " + errorMessage); + } + + @Test + void testConverterRestClientServerErrorPropagates() { + Converter restErrorConverter = (topic, buffer) -> { + RestClientException restException = new RestClientException("server", 500, 50001); + throw new RuntimeException("wrapper", restException); + }; + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), restErrorConverter); + Record kafkaRecord = createKafkaRecord("key".getBytes(StandardCharsets.UTF_8), "value".getBytes(StandardCharsets.UTF_8), null); + + assertThrows(SchemaRegistrySystemException.class, () -> processor.process(TEST_PARTITION, kafkaRecord)); + } + + @Test + void testProcessWithNullRecordThrows() { + DefaultRecordProcessor processor = new DefaultRecordProcessor(TEST_TOPIC, new RawConverter(), new RawConverter()); + + assertThrows(NullPointerException.class, () -> processor.process(TEST_PARTITION, null)); + } + + private Cache extractValueWrapperSchemaCache(DefaultRecordProcessor processor) { + try { + Field cacheField = DefaultRecordProcessor.class.getDeclaredField("valueWrapperSchemaCache"); + cacheField.setAccessible(true); + @SuppressWarnings("unchecked") + Cache cache = (Cache) cacheField.get(processor); + return cache; + } catch (ReflectiveOperationException e) { + throw new AssertionError("Failed to access valueWrapperSchemaCache", e); + } + } + + /** + * Test helper implementations. + */ + private static class SimpleRecord implements Record { + private final long offset; + private final long timestamp; + private final byte[] key; + private final byte[] value; + private final Header[] headers; + + public SimpleRecord(long offset, long timestamp, byte[] key, byte[] value, Header[] headers) { + this.offset = offset; + this.timestamp = timestamp; + this.key = key; + this.value = value; + this.headers = headers != null ? headers : new Header[0]; + } + + @Override + public long offset() { + return offset; + } + + @Override + public int sequence() { + return -1; + } + + @Override + public int sizeInBytes() { + int size = 0; + if (key != null) size += key.length; + if (value != null) size += value.length; + return size; + } + + @Override + public long timestamp() { + return timestamp; + } + + @Override + public void ensureValid() {} + + @Override + public int keySize() { + return key != null ? key.length : -1; + } + + @Override + public boolean hasKey() { + return key != null; + } + + @Override + public ByteBuffer key() { + return key == null ? null : ByteBuffer.wrap(key); + } + + @Override + public int valueSize() { + return value != null ? value.length : -1; + } + + @Override + public boolean hasValue() { + return value != null; + } + + @Override + public ByteBuffer value() { + return value == null ? null : ByteBuffer.wrap(value); + } + + @Override + public boolean hasMagic(byte b) { + return false; + } + + @Override + public boolean isCompressed() { + return false; + } + + @Override + public boolean hasTimestampType(TimestampType timestampType) { + return false; + } + + @Override + public Header[] headers() { + return headers; + } + } + + private static final class NamedPassthroughTransform implements Transform { + private final String name; + + private NamedPassthroughTransform(String name) { + this.name = name; + } + + @Override + public void configure(Map configs) { + // no-op + } + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) { + return record; + } + + @Override + public String getName() { + return name; + } + } + + private static final class NullingTransform implements Transform { + @Override + public void configure(Map configs) { + // no-op + } + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) { + return null; + } + } + + private static final class InvalidDataThrowingTransform implements Transform { + @Override + public void configure(Map configs) { + // no-op + } + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) { + throw new InvalidDataException("Invalid data from transform"); + } + + @Override + public String getName() { + return "InvalidDataTransform"; + } + } + + private static final class ThrowingTransform implements Transform { + @Override + public void configure(Map configs) { + // no-op + } + + @Override + public GenericRecord apply(GenericRecord record, TransformContext context) { + throw new TransformException("transform failure"); + } + } +} diff --git a/core/src/test/java/kafka/automq/table/process/RecordAssemblerTest.java b/core/src/test/java/kafka/automq/table/process/RecordAssemblerTest.java new file mode 100644 index 0000000000..de3fa1470a --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/RecordAssemblerTest.java @@ -0,0 +1,171 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +@Tag("S3Unit") +class RecordAssemblerTest { + + @Test + void ensureOptionalShouldWrapNonUnionSchema() { + Schema stringSchema = Schema.create(Schema.Type.STRING); + + Schema optionalSchema = RecordAssembler.ensureOptional(stringSchema); + + assertEquals(Schema.Type.UNION, optionalSchema.getType()); + List types = optionalSchema.getTypes(); + assertEquals(2, types.size()); + assertEquals(Schema.Type.NULL, types.get(0).getType()); + assertSame(stringSchema, types.get(1)); + } + + @Test + void ensureOptionalShouldPrefixNullWhenMissing() { + Schema stringSchema = Schema.create(Schema.Type.STRING); + Schema intSchema = Schema.create(Schema.Type.INT); + Schema unionWithoutNull = Schema.createUnion(List.of(stringSchema, intSchema)); + + Schema optionalSchema = RecordAssembler.ensureOptional(unionWithoutNull); + + assertEquals(Schema.Type.UNION, optionalSchema.getType()); + List types = optionalSchema.getTypes(); + assertEquals(3, types.size()); + assertEquals(Schema.Type.NULL, types.get(0).getType()); + assertSame(stringSchema, types.get(1)); + assertSame(intSchema, types.get(2)); + } + + @Test + void ensureOptionalShouldReturnOriginalUnionWhenNullPresent() { + Schema unionWithNull = Schema.createUnion(List.of(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + + Schema result = RecordAssembler.ensureOptional(unionWithNull); + + assertSame(unionWithNull, result); + } + + @Test + void assembleShouldExposeOptionalSyntheticFields() { + Schema baseSchema = SchemaBuilder.record("BaseRecord") + .namespace("kafka.automq.table.process.test") + .fields() + .name("name").type().stringType().noDefault() + .endRecord(); + GenericRecord baseRecord = new GenericData.Record(baseSchema); + baseRecord.put("name", "Alice"); + + Schema headerSchema = SchemaBuilder.record("HeaderRecord") + .namespace("kafka.automq.table.process.test") + .fields() + .name("headerKey").type().stringType().noDefault() + .endRecord(); + GenericRecord headerRecord = new GenericData.Record(headerSchema); + headerRecord.put("headerKey", "headerValue"); + ConversionResult headerResult = new ConversionResult(headerRecord, "header-identity"); + + Schema keySchema = Schema.create(Schema.Type.BYTES); + ByteBuffer keyValue = ByteBuffer.wrap("key-value".getBytes(StandardCharsets.UTF_8)); + ConversionResult keyResult = new ConversionResult(keyValue, keySchema, "key-identity"); + + int partition = 5; + long offset = 42L; + long timestamp = 1_700_000_000_000L; + + RecordAssembler assembler = new RecordAssembler(); + GenericRecord assembledRecord = assembler.reset(baseRecord) + .withHeader(headerResult) + .withKey(keyResult) + .withMetadata(partition, offset, timestamp) + .assemble(); + + Schema finalSchema = assembledRecord.getSchema(); + assertEquals(4, finalSchema.getFields().size()); + + Schema headerFieldSchema = finalSchema.getField(RecordAssembler.KAFKA_HEADER_FIELD).schema(); + assertEquals(Schema.Type.UNION, headerFieldSchema.getType()); + assertEquals(Schema.Type.NULL, headerFieldSchema.getTypes().get(0).getType()); + assertSame(headerSchema, headerFieldSchema.getTypes().get(1)); + + Schema keyFieldSchema = finalSchema.getField(RecordAssembler.KAFKA_KEY_FIELD).schema(); + assertEquals(Schema.Type.UNION, keyFieldSchema.getType()); + assertEquals(Schema.Type.NULL, keyFieldSchema.getTypes().get(0).getType()); + assertSame(keySchema, keyFieldSchema.getTypes().get(1)); + + Schema metadataFieldSchema = finalSchema.getField(RecordAssembler.KAFKA_METADATA_FIELD).schema(); + assertEquals(Schema.Type.UNION, metadataFieldSchema.getType()); + assertEquals(Schema.Type.NULL, metadataFieldSchema.getTypes().get(0).getType()); + Schema metadataSchema = metadataFieldSchema.getTypes().get(1); + assertEquals("KafkaMetadata", metadataSchema.getName()); + + assertEquals("Alice", assembledRecord.get("name").toString()); + assertSame(headerRecord, assembledRecord.get(RecordAssembler.KAFKA_HEADER_FIELD)); + assertSame(keyValue, assembledRecord.get(RecordAssembler.KAFKA_KEY_FIELD)); + + GenericRecord metadataRecord = (GenericRecord) assembledRecord.get(RecordAssembler.KAFKA_METADATA_FIELD); + assertNotNull(metadataRecord); + assertEquals(partition, metadataRecord.get(RecordAssembler.METADATA_PARTITION_FIELD)); + assertEquals(offset, metadataRecord.get(RecordAssembler.METADATA_OFFSET_FIELD)); + assertEquals(timestamp, metadataRecord.get(RecordAssembler.METADATA_TIMESTAMP_FIELD)); + } + + @Test + void assembleShouldSkipHeaderWhenAbsent() { + Schema baseSchema = SchemaBuilder.record("BaseRecordNoHeader") + .namespace("kafka.automq.table.process.test") + .fields() + .name("id").type().longType().noDefault() + .endRecord(); + GenericRecord baseRecord = new GenericData.Record(baseSchema); + baseRecord.put("id", 100L); + + Schema keySchema = Schema.create(Schema.Type.STRING); + ConversionResult keyResult = new ConversionResult("primary-key", keySchema, "key-identity"); + + RecordAssembler assembler = new RecordAssembler(); + GenericRecord assembledRecord = assembler.reset(baseRecord) + .withKey(keyResult) + .withMetadata(1, 2L, 3L) + .assemble(); + + Schema finalSchema = assembledRecord.getSchema(); + assertNull(finalSchema.getField(RecordAssembler.KAFKA_HEADER_FIELD)); + assertNotNull(finalSchema.getField(RecordAssembler.KAFKA_KEY_FIELD)); + assertNotNull(finalSchema.getField(RecordAssembler.KAFKA_METADATA_FIELD)); + + assertEquals("primary-key", assembledRecord.get(RecordAssembler.KAFKA_KEY_FIELD)); + GenericRecord metadataRecord = (GenericRecord) assembledRecord.get(RecordAssembler.KAFKA_METADATA_FIELD); + assertNotNull(metadataRecord.getSchema().getField(RecordAssembler.METADATA_PARTITION_FIELD)); + } +} diff --git a/core/src/test/java/kafka/automq/table/process/RecordProcessorFactoryTest.java b/core/src/test/java/kafka/automq/table/process/RecordProcessorFactoryTest.java new file mode 100644 index 0000000000..afa4cdc21f --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/RecordProcessorFactoryTest.java @@ -0,0 +1,922 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process; + +import kafka.automq.table.deserializer.proto.CustomProtobufSchema; +import kafka.automq.table.deserializer.proto.ProtobufSchemaProvider; +import kafka.automq.table.process.exception.ProcessorInitializationException; +import kafka.automq.table.process.exception.TransformException; +import kafka.automq.table.process.proto.PersonProto; +import kafka.automq.table.worker.WorkerConfig; + +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.header.internals.RecordHeader; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.server.record.TableTopicConvertType; +import org.apache.kafka.server.record.TableTopicTransformType; + +import com.google.protobuf.Timestamp; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; + +import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; +import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; +import io.confluent.kafka.serializers.KafkaAvroSerializer; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +public class RecordProcessorFactoryTest { + + private static final String TEST_TOPIC = "test-topic"; + private static final int TEST_PARTITION = 0; + private static final long TEST_OFFSET = 123L; + private static final long TEST_TIMESTAMP = System.currentTimeMillis(); + + private static final Schema USER_SCHEMA = SchemaBuilder.record("User") + .namespace("kafka.automq.table.process") + .fields() + .name("name").type().stringType().noDefault() + .name("age").type().intType().noDefault() + .endRecord(); + + private static final Schema PRODUCT_SCHEMA = SchemaBuilder.record("Product") + .namespace("kafka.automq.table.process") + .fields() + .name("product_id").type().longType().noDefault() + .name("product_name").type().stringType().noDefault() + .name("price").type().doubleType().noDefault() + .endRecord(); + + private static final Schema DEBEZIUM_ENVELOPE_SCHEMA = SchemaBuilder.record("Envelope") + .namespace("io.debezium.connector.mysql") + .fields() + .name("before").type().unionOf().nullType().and().type(PRODUCT_SCHEMA).endUnion().noDefault() + .name("after").type().unionOf().nullType().and().type(PRODUCT_SCHEMA).endUnion().noDefault() + .name("source").type(SchemaBuilder.record("Source") + .fields() + .name("db").type().stringType().noDefault() + .name("table").type().stringType().noDefault() + .endRecord()) + .noDefault() + .name("op").type().stringType().noDefault() + .name("ts_ms").type().longType().noDefault() + .endRecord(); + + private SchemaRegistryClient schemaRegistryClient; + private RecordProcessorFactory recordProcessorFactory; + private KafkaAvroSerializer avroSerializer; + + @BeforeEach + void setUp() { + schemaRegistryClient = new MockSchemaRegistryClient(List.of(new ProtobufSchemaProvider())); + recordProcessorFactory = new RecordProcessorFactory("http://mock:8081", schemaRegistryClient); + avroSerializer = new KafkaAvroSerializer(schemaRegistryClient); + avroSerializer.configure(Map.of("schema.registry.url", "http://mock:8081"), false); + } + + // --- Test Group 1: RAW Converter --- + + @Test + void testRawConverterWithoutTransforms() { + // Arrange + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + byte[] key = "test-key".getBytes(); + byte[] value = "test-value".getBytes(); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, value, key); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess(), "Processing should be successful"); + assertNull(result.getError(), "Error should be null"); + + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check the new standard output structure + assertTrue(finalRecord.hasField("_kafka_key")); + assertTrue(finalRecord.hasField("_kafka_value")); + assertTrue(finalRecord.hasField("_kafka_header")); + assertTrue(finalRecord.hasField("_kafka_metadata")); + + assertEquals(ByteBuffer.wrap(key), ByteBuffer.wrap((byte[]) finalRecord.get("_kafka_key"))); + assertEquals(ByteBuffer.wrap(value), ByteBuffer.wrap((byte[]) finalRecord.get("_kafka_value"))); + + GenericRecord metadataRecord = (GenericRecord) finalRecord.get("_kafka_metadata"); + assertEquals(TEST_PARTITION, metadataRecord.get("partition")); + assertEquals(TEST_OFFSET, metadataRecord.get("offset")); + assertEquals(TEST_TIMESTAMP, metadataRecord.get("timestamp")); + } + + // --- Test Group 2: BY_SCHEMA_ID Converter --- + + @Test + void testBySchemaIdWithAvro() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + Schema schema = Schema.create(Schema.Type.STRING); + int schemaId = registerSchema(subject, schema); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + // Create an actual Avro record with the string value + String avroStringValue = "test123"; + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, avroStringValue, "test-key"); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + + + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check the new standard output structure + assertTrue(finalRecord.hasField("_kafka_key")); + assertTrue(finalRecord.hasField("_kafka_value")); + assertTrue(finalRecord.hasField("_kafka_metadata")); + + assertEquals("test-key", finalRecord.get("_kafka_key")); + assertEquals("test123", finalRecord.get("_kafka_value")); + + GenericRecord metadataRecord = (GenericRecord) finalRecord.get("_kafka_metadata"); + assertEquals(TEST_PARTITION, metadataRecord.get("partition")); + assertEquals(TEST_OFFSET, metadataRecord.get("offset")); + assertEquals(TEST_TIMESTAMP, metadataRecord.get("timestamp")); + } + + @Test + void testBySchemaIdWithUnwrap() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + int schemaId = registerSchema(subject, USER_SCHEMA); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + GenericRecord userRecord = new GenericRecordBuilder(USER_SCHEMA) + .set("name", "test-user") + .set("age", 30) + .build(); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, userRecord, "test-key"); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + + + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check that flatten has been applied - the value fields should be at the top level + assertTrue(finalRecord.hasField("name")); + assertTrue(finalRecord.hasField("age")); + assertEquals("test-user", finalRecord.get("name").toString()); + assertEquals(30, finalRecord.get("age")); + + // Check that kafka fields are still present + assertTrue(finalRecord.hasField("_kafka_key")); + assertTrue(finalRecord.hasField("_kafka_metadata")); + assertEquals("test-key", finalRecord.get("_kafka_key")); + + GenericRecord metadataRecord = (GenericRecord) finalRecord.get("_kafka_metadata"); + assertEquals(TEST_PARTITION, metadataRecord.get("partition")); + assertEquals(TEST_OFFSET, metadataRecord.get("offset")); + assertEquals(TEST_TIMESTAMP, metadataRecord.get("timestamp")); + } + + // --- Test Group 3: Debezium Unwrap --- + + @Test + void testBySchemaIdWithDebeziumUnwrap() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + int schemaId = registerSchema(subject, DEBEZIUM_ENVELOPE_SCHEMA); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN_DEBEZIUM); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + GenericRecord productRecord = new GenericRecordBuilder(PRODUCT_SCHEMA) + .set("product_id", 1L) + .set("product_name", "test-product") + .set("price", 99.99) + .build(); + + GenericRecord sourceRecord = new GenericRecordBuilder(DEBEZIUM_ENVELOPE_SCHEMA.getField("source").schema()) + .set("db", "test_db") + .set("table", "products") + .build(); + + GenericRecord debeziumRecord = new GenericRecordBuilder(DEBEZIUM_ENVELOPE_SCHEMA) + .set("before", null) + .set("after", productRecord) + .set("source", sourceRecord) + .set("op", "c") + .set("ts_ms", TEST_TIMESTAMP) + .build(); + + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, debeziumRecord, "test-key"); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check that the product fields are flattened to top level + assertEquals(1L, finalRecord.get("product_id")); + assertEquals("test-product", finalRecord.get("product_name").toString()); + assertEquals(99.99, (Double) finalRecord.get("price"), 0.01); + + // Check CDC metadata + assertTrue(finalRecord.hasField("_cdc")); + GenericRecord cdcRecord = (GenericRecord) finalRecord.get("_cdc"); + assertEquals("I", cdcRecord.get("op").toString()); + assertEquals(TEST_TIMESTAMP, cdcRecord.get("ts")); + assertEquals(TEST_OFFSET, cdcRecord.get("offset")); + assertEquals("test_db.products", cdcRecord.get("source").toString()); + + // Check Kafka metadata + assertTrue(finalRecord.hasField("_kafka_key")); + assertTrue(finalRecord.hasField("_kafka_metadata")); + assertEquals("test-key", finalRecord.get("_kafka_key")); + + GenericRecord metadataRecord = (GenericRecord) finalRecord.get("_kafka_metadata"); + assertEquals(TEST_PARTITION, metadataRecord.get("partition")); + assertEquals(TEST_OFFSET, metadataRecord.get("offset")); + assertEquals(TEST_TIMESTAMP, metadataRecord.get("timestamp")); + } + + // --- Test Group 4: BY_SUBJECT_NAME Converter --- + + @Test + void testBySubjectNameThrowsExceptionForNonProtobufSchema() throws Exception { + // Arrange + String subject = "avro-subject"; + registerSchema(subject, USER_SCHEMA); // Register an AVRO schema + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_LATEST_SCHEMA); + when(mockConfig.valueSubject()).thenReturn(subject); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, "any-topic"); + Record kafkaRecord = createKafkaRecord("any-topic", "dummy-payload".getBytes(), "test-key"); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertFalse(result.isSuccess()); + assertNotNull(result.getError()); + assertTrue(result.getError().getCause() instanceof ProcessorInitializationException, "Cause should be ProcessorInitializationException"); + assertTrue(result.getError().getCause().getMessage().contains("by_subject_name is only supported for PROTOBUF")); + } + + @Test + void testBySubjectNameWithProtobufSchema() throws Exception { + // Arrange + String subject = "proto-person-subject"; + String protoFileContent = Files.readString(Path.of("src/test/resources/proto/person.proto")); + + CustomProtobufSchema person = new CustomProtobufSchema("Person", -1, null, null, protoFileContent, List.of(), Map.of()); + int schemaId = schemaRegistryClient.register(subject, person); + + String messageFullName = "kafka.automq.table.process.proto.Person"; + + PersonProto.Address address = PersonProto.Address.newBuilder() + .setStreet("123 Main St") + .setCity("Anytown") + .build(); + + long now = System.currentTimeMillis(); + Timestamp timestamp = Timestamp.newBuilder().setSeconds(now / 1000).setNanos((int) ((now % 1000) * 1000000)).build(); + + PersonProto.Person personMessage = PersonProto.Person.newBuilder() + .setId(1L) + .setName("Proto User") + .setIsActive(true) + .setAddress(address) + .addRoles("admin") + .addRoles("user") + .putAttributes("team", "backend") + .setLastUpdated(timestamp) + .build(); + + byte[] value = personMessage.toByteArray(); + Record kafkaRecord = createKafkaRecord(subject, value, "test-key".getBytes()); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_LATEST_SCHEMA); + when(mockConfig.valueSubject()).thenReturn(subject); + when(mockConfig.valueMessageFullName()).thenReturn(messageFullName); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, subject); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check that flatten has been applied - the protobuf fields should be at the top level + assertEquals(1L, finalRecord.get("id")); + assertEquals("Proto User", finalRecord.get("name").toString()); + assertEquals(true, finalRecord.get("is_active")); + + // Check nested record + GenericRecord addressRecord = (GenericRecord) finalRecord.get("address"); + assertEquals("123 Main St", addressRecord.get("street").toString()); + + // Check repeated field (list) + @SuppressWarnings("unchecked") + List roles = (List) finalRecord.get("roles"); + assertEquals(2, roles.size()); + assertEquals("admin", roles.get(0).toString()); + + // Check map (which is converted to a list of records) + @SuppressWarnings("unchecked") + List attributesList = (List) finalRecord.get("attributes"); + assertNotNull(attributesList); + assertEquals(1, attributesList.size()); + GenericRecord attributeEntry = attributesList.get(0); + assertEquals("team", attributeEntry.get("key").toString()); + assertEquals("backend", attributeEntry.get("value").toString()); + + // Check timestamp (converted to long in microseconds) + long expectedTimestampMicros = timestamp.getSeconds() * 1_000_000 + timestamp.getNanos() / 1000; + assertEquals(expectedTimestampMicros, finalRecord.get("last_updated")); + + // Check Kafka fields + assertTrue(finalRecord.hasField("_kafka_key")); + assertTrue(finalRecord.hasField("_kafka_metadata")); + assertEquals("test-key", finalRecord.get("_kafka_key")); + + GenericRecord metadataRecord = (GenericRecord) finalRecord.get("_kafka_metadata"); + assertEquals(TEST_PARTITION, metadataRecord.get("partition")); + assertEquals(TEST_OFFSET, metadataRecord.get("offset")); + assertEquals(TEST_TIMESTAMP, metadataRecord.get("timestamp")); + } + + + @Test + void testBySubjectNameWithFirstProtobufSchema() throws Exception { + // Arrange + String subject = "proto-address-subject"; + String protoFileContent = Files.readString(Path.of("src/test/resources/proto/person.proto")); + + CustomProtobufSchema addressSchema = new CustomProtobufSchema("Address", -1, null, null, protoFileContent, List.of(), Map.of()); + int schemaId = schemaRegistryClient.register(subject, addressSchema); + + String messageFullName = null; + + PersonProto.Address address = PersonProto.Address.newBuilder() + .setStreet("123 Main St") + .setCity("Anytown") + .build(); + + byte[] value = address.toByteArray(); + Record kafkaRecord = createKafkaRecord(subject, value, "test-key".getBytes()); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_LATEST_SCHEMA); + when(mockConfig.valueSubject()).thenReturn(subject); + when(mockConfig.valueMessageFullName()).thenReturn(messageFullName); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, subject); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + + GenericRecord addressRecord = result.getFinalRecord(); + assertEquals("123 Main St", addressRecord.get("street").toString()); + } + + // --- Test Group 5: Error Handling --- + + @Test + void testConvertErrorOnUnknownSchemaId() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + registerSchema(subject, USER_SCHEMA); + GenericRecord userRecord = new GenericRecordBuilder(USER_SCHEMA).set("name", "a").set("age", 1).build(); + byte[] validPayload = avroSerializer.serialize(TEST_TOPIC, userRecord); + + ByteBuffer buffer = ByteBuffer.wrap(validPayload); + buffer.get(); // Magic byte + buffer.putInt(9999); // Non-existent schema ID + byte[] invalidPayload = buffer.array(); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, invalidPayload, "test-key".getBytes()); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertFalse(result.isSuccess()); + assertNotNull(result.getError()); + assertEquals(DataError.ErrorType.CONVERT_ERROR, result.getError().getType()); + assertTrue(result.getError().getCause() instanceof RestClientException); + } + + @Test + void testTransformErrorOnMismatchedData() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + registerSchema(subject, USER_SCHEMA); + GenericRecord userRecord = new GenericRecordBuilder(USER_SCHEMA).set("name", "a").set("age", 1).build(); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, userRecord, "test-key"); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN_DEBEZIUM); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertFalse(result.isSuccess()); + assertNotNull(result.getError()); + assertEquals(DataError.ErrorType.TRANSFORMATION_ERROR, result.getError().getType()); + assertTrue(result.getError().getCause() instanceof TransformException); + } + + + @Test + void testHeaderConversion() { + // Arrange + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + // Create headers + Header[] headers = new Header[] { + new RecordHeader("content-type", "application/json".getBytes()), + new RecordHeader("source-system", "order-service".getBytes()) + }; + + byte[] key = "test-key".getBytes(); + byte[] value = "test-value".getBytes(); + Record kafkaRecord = new SimpleRecord(TEST_OFFSET, TEST_TIMESTAMP, key, value, headers); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess(), "Processing should be successful"); + assertNull(result.getError(), "Error should be null"); + + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check headers + assertTrue(finalRecord.hasField("_kafka_header")); + @SuppressWarnings("unchecked") + Map headerMap = (Map) finalRecord.get("_kafka_header"); + assertNotNull(headerMap); + assertEquals(2, headerMap.size()); + assertEquals("application/json", new String(headerMap.get("content-type").array())); + assertEquals("order-service", new String(headerMap.get("source-system").array())); + } + + @Test + void testKeyConversionWithSchema() throws Exception { + // Arrange + String keySubject = TEST_TOPIC + "-key"; + String valueSubject = TEST_TOPIC + "-value"; + + Schema keySchema = Schema.create(Schema.Type.STRING); + registerSchema(keySubject, keySchema); + registerSchema(valueSubject, USER_SCHEMA); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.keySubject()).thenReturn(keySubject); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + // Create Avro serialized key and raw value + String keyData = "user-123"; + byte[] serializedKey = avroSerializer.serialize(keySubject, keyData); + byte[] rawValue = "raw-value-data".getBytes(); + + Record kafkaRecord = new SimpleRecord(TEST_OFFSET, TEST_TIMESTAMP, serializedKey, rawValue); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check key is converted to GenericRecord + assertTrue(finalRecord.hasField("_kafka_key")); + assertEquals("user-123", finalRecord.get("_kafka_key")); + + // Check value remains as ByteBuffer + assertTrue(finalRecord.hasField("_kafka_value")); + assertEquals(ByteBuffer.wrap(rawValue), ByteBuffer.wrap((byte[]) finalRecord.get("_kafka_value"))); + } + + @Test + void testKeyConversionAsString() { + // Arrange + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + byte[] keyBytes = "user-456".getBytes(); + byte[] value = "test-value".getBytes(); + Record kafkaRecord = new SimpleRecord(TEST_OFFSET, TEST_TIMESTAMP, keyBytes, value); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check key is converted to String + assertTrue(finalRecord.hasField("_kafka_key")); + assertEquals("user-456", finalRecord.get("_kafka_key")); + + // Check value remains as ByteBuffer + assertTrue(finalRecord.hasField("_kafka_value")); + assertEquals(ByteBuffer.wrap(value), ByteBuffer.wrap((byte[]) finalRecord.get("_kafka_value"))); + } + + // --- Test Group 6: Deprecated Configs --- + + @Test + void testSchemalessConfig() { + // Arrange + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.schemaType()).thenReturn(org.apache.kafka.server.record.TableTopicSchemaType.SCHEMALESS); + // These should be ignored when schemaType is SCHEMALESS + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + String key = "schemaless-key"; + String value = "schemaless-value"; + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, value.getBytes(StandardCharsets.UTF_8), key.getBytes(StandardCharsets.UTF_8)); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check that the SchemalessTransform was applied + assertTrue(finalRecord.hasField("key")); + assertTrue(finalRecord.hasField("value")); + assertTrue(finalRecord.hasField("timestamp")); + + assertEquals(key, finalRecord.get("key")); + assertEquals(value, finalRecord.get("value")); + assertEquals(TEST_TIMESTAMP, finalRecord.get("timestamp")); + + // Check that Kafka metadata is still present + assertTrue(finalRecord.hasField("_kafka_key")); + assertTrue(finalRecord.hasField("_kafka_metadata")); + assertTrue(finalRecord.hasField("_kafka_header")); + assertFalse(finalRecord.hasField("_kafka_value"), + "Value should be unwrapped by SchemalessTransform, not present as _kafka_value"); + + // The key converter in schemaless mode is RawConverter + assertEquals(key, finalRecord.get("_kafka_key")); + } + + @Test + void testSchemaConfigDeprecated() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + registerSchema(subject, USER_SCHEMA); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.schemaType()).thenReturn(org.apache.kafka.server.record.TableTopicSchemaType.SCHEMA); + // This should be used by the value converter + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + // These should be ignored by the factory logic for this deprecated config + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + + GenericRecord userRecord = new GenericRecordBuilder(USER_SCHEMA) + .set("name", "test-user-deprecated") + .set("age", 40) + .build(); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, userRecord, "test-key-deprecated"); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess()); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check that FlattenTransform was applied (as per deprecated config logic) + assertEquals("test-user-deprecated", finalRecord.get("name").toString()); + assertEquals(40, finalRecord.get("age")); + + // Check that key was processed by StringConverter (as per deprecated config logic) + assertEquals("test-key-deprecated", finalRecord.get("_kafka_key")); + } + + // --- Test Group 7: More Converter/Error Scenarios --- + + @Test + void testBySchemaIdWithProtobuf() throws Exception { + // Arrange + String subject = TEST_TOPIC + "-value"; + String protoFileContent = Files.readString(Path.of("src/test/resources/proto/person.proto")); + CustomProtobufSchema person = new CustomProtobufSchema("Person", -1, null, null, protoFileContent, List.of(), Map.of()); + int schemaId = schemaRegistryClient.register(subject, person); + + PersonProto.Address addressMessage = PersonProto.Address.newBuilder() + .setStreet("456 Oak Ave") + .setCity("Othertown") + .build(); + + // Manually construct the payload with magic byte, schema ID, and message index. + // Based on the Confluent wire format, an index array size of 0 defaults to the first message type ([0]). + // This is a compact representation for the most common case. + byte[] protoBytes = addressMessage.toByteArray(); + byte[] messageIndexArray = new byte[]{0}; // A size of 0 defaults to index [0] + ByteBuffer buffer = ByteBuffer.allocate(1 + 4 + messageIndexArray.length + protoBytes.length); + buffer.put((byte) 0x0); + buffer.putInt(schemaId); + buffer.put(messageIndexArray); + buffer.put(protoBytes); + byte[] valuePayload = buffer.array(); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.STRING); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.FLATTEN); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, valuePayload, "proto-key".getBytes()); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertTrue(result.isSuccess(), "Processing should be successful"); + GenericRecord finalRecord = result.getFinalRecord(); + assertNotNull(finalRecord); + + // Check flattened fields from Protobuf + assertEquals("456 Oak Ave", finalRecord.get("street").toString()); + assertEquals("Othertown", finalRecord.get("city").toString()); + assertFalse(finalRecord.hasField("id"), "Fields from Person message should not be present"); + + // Check Kafka metadata + assertEquals("proto-key", finalRecord.get("_kafka_key")); + } + + @Test + void testKeyConvertError() throws Exception { + // Arrange + String keySubject = TEST_TOPIC + "-key"; + registerSchema(keySubject, Schema.create(Schema.Type.STRING)); + + // Create a key payload with an invalid schema ID + ByteBuffer buffer = ByteBuffer.allocate(5); + buffer.put((byte) 0x0); + buffer.putInt(9999); // Invalid ID + byte[] invalidKey = buffer.array(); + + WorkerConfig mockConfig = mock(WorkerConfig.class); + when(mockConfig.keyConvertType()).thenReturn(TableTopicConvertType.BY_SCHEMA_ID); + when(mockConfig.valueConvertType()).thenReturn(TableTopicConvertType.RAW); + when(mockConfig.transformType()).thenReturn(TableTopicTransformType.NONE); + + RecordProcessor processor = recordProcessorFactory.create(mockConfig, TEST_TOPIC); + Record kafkaRecord = createKafkaRecord(TEST_TOPIC, "some-value".getBytes(), invalidKey); + + // Act + ProcessingResult result = processor.process(TEST_PARTITION, kafkaRecord); + + // Assert + assertFalse(result.isSuccess()); + assertNotNull(result.getError()); + assertEquals(DataError.ErrorType.CONVERT_ERROR, result.getError().getType()); + assertTrue(result.getError().getCause() instanceof RestClientException, + "Cause should be RestClientException from the deserializer"); + } + + // --- Helper Methods --- + + private Record createKafkaRecord(String topic, byte[] value, byte[] key) { + return new SimpleRecord(TEST_OFFSET, TEST_TIMESTAMP, key, value); + } + + private Record createKafkaRecord(String topic, Object avroRecord, String key) { + byte[] value = avroSerializer.serialize(topic, avroRecord); + byte[] keyBytes = key.getBytes(); + return createKafkaRecord(topic, value, keyBytes); + } + + private int registerSchema(String subject, Schema schema) throws Exception { + return schemaRegistryClient.register(subject, new io.confluent.kafka.schemaregistry.avro.AvroSchema(schema)); + } + + /** + * A simplified implementation of the Record interface for testing purposes. + */ + private static class SimpleRecord implements Record { + private final long offset; + private final long timestamp; + private final byte[] key; + private final byte[] value; + private final Header[] headers; + + public SimpleRecord(long offset, long timestamp, byte[] key, byte[] value) { + this(offset, timestamp, key, value, new Header[0]); + } + + public SimpleRecord(long offset, long timestamp, byte[] key, byte[] value, Header[] headers) { + this.offset = offset; + this.timestamp = timestamp; + this.key = key; + this.value = value; + this.headers = headers != null ? headers : new Header[0]; + } + + @Override + public long offset() { + return offset; + } + + @Override + public int sequence() { + return -1; + } + + @Override + public int sizeInBytes() { + int size = 0; + if (key != null) size += key.length; + if (value != null) size += value.length; + return size; + } + + @Override + public long timestamp() { + return timestamp; + } + + @Override + public void ensureValid() {} + + @Override + public int keySize() { + return key != null ? key.length : -1; + } + + @Override + public boolean hasKey() { + return key != null; + } + + @Override + public ByteBuffer key() { + return key == null ? null : ByteBuffer.wrap(key); + } + + @Override + public int valueSize() { + return value != null ? value.length : -1; + } + + @Override + public boolean hasValue() { + return value != null; + } + + @Override + public ByteBuffer value() { + return value == null ? null : ByteBuffer.wrap(value); + } + + @Override + public boolean hasMagic(byte b) { + return false; + } + + @Override + public boolean isCompressed() { + return false; + } + + @Override + public boolean hasTimestampType(TimestampType timestampType) { + return false; + } + + @Override + public Header[] headers() { + return headers; + } + } +} diff --git a/core/src/test/java/kafka/automq/table/process/SchemalessTransformTest.java b/core/src/test/java/kafka/automq/table/process/SchemalessTransformTest.java new file mode 100644 index 0000000000..eceef5fadb --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/SchemalessTransformTest.java @@ -0,0 +1,109 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process; + +import kafka.automq.table.process.exception.TransformException; +import kafka.automq.table.process.transform.SchemalessTransform; + +import org.apache.kafka.common.record.Record; + +import org.apache.avro.generic.GenericRecord; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import java.nio.ByteBuffer; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +public class SchemalessTransformTest { + + @Mock + private Record kafkaRecord; + + @Mock + private GenericRecord inputRecord; + + @Mock + private TransformContext context; + + private SchemalessTransform transform; + + @BeforeEach + public void setup() { + MockitoAnnotations.openMocks(this); + transform = new SchemalessTransform(); + when(context.getKafkaRecord()).thenReturn(kafkaRecord); + } + + @Test + public void testApply_WithKeyAndValue() throws TransformException { + // Arrange + long timestamp = 1234567890L; + + when(kafkaRecord.hasKey()).thenReturn(true); + when(kafkaRecord.key()).thenReturn(ByteBuffer.wrap("321".getBytes())); + when(kafkaRecord.hasValue()).thenReturn(true); + when(kafkaRecord.value()).thenReturn(ByteBuffer.wrap("123".getBytes())); + when(kafkaRecord.timestamp()).thenReturn(timestamp); + + // Act + GenericRecord result = transform.apply(inputRecord, context); + + // Assert + assertNotNull(result); + assertEquals(SchemalessTransform.SCHEMALESS_SCHEMA, result.getSchema()); + assertEquals("321", (String) result.get("key")); + assertEquals("123", (String) result.get("value")); + assertEquals(timestamp, result.get("timestamp")); + } + + @Test + public void testApply_WithoutKeyAndValue() throws TransformException { + // Arrange + long timestamp = 1234567890L; + + when(kafkaRecord.hasKey()).thenReturn(false); + when(kafkaRecord.hasValue()).thenReturn(false); + when(kafkaRecord.timestamp()).thenReturn(timestamp); + + // Act + GenericRecord result = transform.apply(inputRecord, context); + + // Assert + assertNotNull(result); + assertEquals(SchemalessTransform.SCHEMALESS_SCHEMA, result.getSchema()); + assertNull(result.get("key")); + assertNull(result.get("value")); + assertEquals(timestamp, result.get("timestamp")); + } + + @Test + public void testGetName() { + // Act & Assert + assertEquals("schemaless", transform.getName()); + } +} diff --git a/core/src/test/java/kafka/automq/table/process/convert/ProtoToAvroConverterTest.java b/core/src/test/java/kafka/automq/table/process/convert/ProtoToAvroConverterTest.java new file mode 100644 index 0000000000..ece4c043c3 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/convert/ProtoToAvroConverterTest.java @@ -0,0 +1,168 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import kafka.automq.table.deserializer.proto.parse.ProtobufSchemaParser; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants; +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; +import kafka.automq.table.process.exception.ConverterException; + +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.protobuf.ProtobufData; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; +import java.util.function.Consumer; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Focused unit tests for {@link ProtoToAvroConverter} exercise converter paths that are + * hard to reach through the higher-level registry converter integration tests. + */ +class ProtoToAvroConverterTest { + + private static final String SIMPLE_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + message SimpleRecord { + bool flag = 1; + Nested nested = 2; + optional int32 opt_scalar = 3; + } + + message Nested { + string note = 1; + } + """; + + @Test + void skipsUnknownAvroFieldsWhenSchemaHasExtraColumns() throws Exception { + DynamicMessage message = buildSimpleRecord(b -> + b.setField(b.getDescriptorForType().findFieldByName("flag"), true) + ); + + Schema schema = SchemaBuilder.record("SimpleRecord") + .fields() + .name("flag").type().booleanType().noDefault() + .name("ghost_field").type().stringType().noDefault() + .endRecord(); + + GenericRecord record = ProtoToAvroConverter.convert(message, schema); + assertEquals(true, record.get("flag")); + assertNull(record.get("ghost_field")); + } + + @Test + void leavesMissingPresenceFieldUnsetWhenAvroSchemaDisallowsNull() throws Exception { + DynamicMessage message = buildSimpleRecord(b -> + b.setField(b.getDescriptorForType().findFieldByName("flag"), false) + ); + + Schema nestedSchema = SchemaBuilder.record("Nested") + .fields() + .name("note").type().stringType().noDefault() + .endRecord(); + + Schema schema = SchemaBuilder.record("SimpleRecord") + .fields() + .name("nested").type(nestedSchema).noDefault() + .name("opt_scalar").type().intType().noDefault() + .endRecord(); + + GenericRecord record = ProtoToAvroConverter.convert(message, schema); + assertNull(record.get("nested")); + assertEquals(0, record.get("opt_scalar")); + } + + @Test + void messageSchemaMismatchYieldsNullWhenNonRecordTypeProvided() throws Exception { + DynamicMessage message = buildSimpleRecord(b -> { + Descriptors.Descriptor nestedDesc = b.getDescriptorForType().findFieldByName("nested").getMessageType(); + b.setField(b.getDescriptorForType().findFieldByName("nested"), + DynamicMessage.newBuilder(nestedDesc) + .setField(nestedDesc.findFieldByName("note"), "note-value") + .build() + ); + }); + + Schema schema = SchemaBuilder.record("SimpleRecord") + .fields() + .name("nested").type().longType().noDefault() + .endRecord(); + + GenericRecord record = ProtoToAvroConverter.convert(message, schema); + assertNull(record.get("nested")); + } + + @Test + void convertPrimitiveWrapsByteArrayValues() throws Exception { + Method method = ProtoToAvroConverter.class.getDeclaredMethod("convertPrimitive", Object.class, Schema.class); + method.setAccessible(true); + byte[] source = new byte[]{1, 2, 3}; + ByteBuffer buffer = (ByteBuffer) invoke(method, null, source, Schema.create(Schema.Type.BYTES)); + ByteBuffer copy = buffer.duplicate(); + byte[] actual = new byte[copy.remaining()]; + copy.get(actual); + assertEquals(List.of((byte) 1, (byte) 2, (byte) 3), List.of(actual[0], actual[1], actual[2])); + } + + @Test + void convertSingleValueRejectsRawListsWhenFieldIsNotRepeated() throws Exception { + Method method = ProtoToAvroConverter.class.getDeclaredMethod("convertSingleValue", Object.class, Schema.class, ProtobufData.class); + method.setAccessible(true); + Schema schema = Schema.create(Schema.Type.STRING); + assertThrows(ConverterException.class, () -> invoke(method, null, List.of("unexpected"), schema, LogicalMapProtobufData.get())); + } + + private static T invoke(Method method, Object target, Object... args) throws Exception { + try { + return (T) method.invoke(target, args); + } catch (InvocationTargetException e) { + throw (Exception) e.getCause(); + } + } + + private static DynamicMessage buildSimpleRecord(Consumer configurer) throws Exception { + Descriptors.Descriptor descriptor = getDescriptor(SIMPLE_PROTO, "SimpleRecord"); + DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptor); + configurer.accept(builder); + return builder.build(); + } + + private static Descriptors.Descriptor getDescriptor(String proto, String messageName) throws Exception { + com.squareup.wire.schema.internal.parser.ProtoFileElement fileElement = + com.squareup.wire.schema.internal.parser.ProtoParser.Companion.parse(ProtoConstants.DEFAULT_LOCATION, proto); + DynamicSchema dynamicSchema = ProtobufSchemaParser.toDynamicSchema(messageName, fileElement, Collections.emptyMap()); + return dynamicSchema.getMessageDescriptor(messageName); + } +} diff --git a/core/src/test/java/kafka/automq/table/process/convert/ProtobufRegistryConverterTest.java b/core/src/test/java/kafka/automq/table/process/convert/ProtobufRegistryConverterTest.java new file mode 100644 index 0000000000..c8b080fdce --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/convert/ProtobufRegistryConverterTest.java @@ -0,0 +1,464 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import kafka.automq.table.binder.RecordBinder; +import kafka.automq.table.deserializer.proto.CustomProtobufSchema; +import kafka.automq.table.deserializer.proto.ProtobufSchemaProvider; +import kafka.automq.table.deserializer.proto.parse.ProtobufSchemaParser; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants; +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; +import kafka.automq.table.process.ConversionResult; + +import org.apache.kafka.common.utils.ByteUtils; + +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Timestamp; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.ProtoParser; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.iceberg.Table; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.TaskWriter; +import org.apache.iceberg.types.Type; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient; + +import static kafka.automq.table.binder.AvroRecordBinderTypeTest.createTableWriter; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +@Tag("S3Unit") +public class ProtobufRegistryConverterTest { + + private static final String ALL_TYPES_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + import \"google/protobuf/timestamp.proto\"; + + message Nested { + string name = 1; + int32 count = 2; + } + + enum SampleEnum { + SAMPLE_ENUM_UNSPECIFIED = 0; + SAMPLE_ENUM_SECOND = 1; + } + + message FloatArray { + repeated double values = 1; + } + + message StringArray { + repeated string values = 1; + } + + message AllTypes { + // Scalar primitives in order defined by Avro ProtobufData mapping + bool f_bool = 1; + double f_double = 2; + float f_float = 3; + int32 f_int32 = 4; + sint32 f_sint32 = 5; + uint32 f_uint32 = 6; + fixed32 f_fixed32 = 7; + sfixed32 f_sfixed32 = 8; + int64 f_int64 = 9; + sint64 f_sint64 = 10; + uint64 f_uint64 = 11; + fixed64 f_fixed64 = 12; + sfixed64 f_sfixed64 = 13; + string f_string = 14; + bytes f_bytes = 15; + SampleEnum f_enum = 16; + Nested f_message = 17; + // Containers and complex types + repeated int32 f_repeated_int32 = 18; + map f_string_int32_map = 19; + google.protobuf.Timestamp f_timestamp = 20; + optional string f_optional_string = 21; + oneof choice { + string choice_str = 22; + int32 choice_int = 23; + FloatArray choice_float_array = 26; + StringArray choice_string_array = 27; + } + repeated Nested f_nested_list = 24; + map f_string_nested_map = 25; + map f_string_timestamp_map = 28; + } + """; + + private static final String MAP_ONLY_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + message MapOnly { + map attributes = 1; + } + """; + + private void testSendRecord(org.apache.iceberg.Schema schema, org.apache.iceberg.data.Record record) { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("test", ImmutableMap.of()); + catalog.createNamespace(Namespace.of("default")); + String tableName = "test"; + Table table = catalog.createTable(TableIdentifier.of(Namespace.of("default"), tableName), schema); + TaskWriter writer = createTableWriter(table); + try { + writer.write(record); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Test + void testConvertAllPrimitiveAndCollectionTypes() throws Exception { + String topic = "pb-all-types"; + String subject = topic + "-value"; + + MockSchemaRegistryClient registryClient = new MockSchemaRegistryClient(List.of(new ProtobufSchemaProvider())); + CustomProtobufSchema schema = new CustomProtobufSchema( + "AllTypes", + -1, + null, + null, + ALL_TYPES_PROTO, + List.of(), + Map.of() + ); + int schemaId = registryClient.register(subject, schema); + + ProtoFileElement fileElement = ProtoParser.Companion.parse(ProtoConstants.DEFAULT_LOCATION, ALL_TYPES_PROTO); + DynamicSchema dynamicSchema = ProtobufSchemaParser.toDynamicSchema("AllTypes", fileElement, Collections.emptyMap()); + Descriptors.Descriptor descriptor = dynamicSchema.getMessageDescriptor("AllTypes"); + + DynamicMessage message = buildAllTypesMessage(descriptor); + // magic byte + schema id + single message index + serialized protobuf payload + ByteBuffer payload = buildConfluentPayload(schemaId, message.toByteArray(), 3); + + ProtobufRegistryConverter converter = new ProtobufRegistryConverter(registryClient, "http://mock:8081", false); + + ConversionResult result = converter.convert(topic, payload.asReadOnlyBuffer()); + ConversionResult cachedResult = converter.convert(topic, buildConfluentPayload(schemaId, message.toByteArray(), 1)); + assertSame(result.getSchema(), cachedResult.getSchema(), "Schema cache should return the same Avro schema instance"); + assertEquals(String.valueOf(schemaId), result.getSchemaIdentity()); + + GenericRecord record = (GenericRecord) result.getValue(); + assertPrimitiveFields(record); + assertRepeatedAndMapFields(record); + assertNestedAndTimestamp(record); + + org.apache.iceberg.Schema iceberg = AvroSchemaUtil.toIceberg(record.getSchema()); + RecordBinder recordBinder = new RecordBinder(iceberg, record.getSchema()); + Record bind = recordBinder.bind(record); + testSendRecord(iceberg, bind); + } + + private static DynamicMessage buildAllTypesMessage(Descriptors.Descriptor descriptor) { + DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptor); + + builder.setField(descriptor.findFieldByName("f_bool"), true); + builder.setField(descriptor.findFieldByName("f_double"), 123.456d); + builder.setField(descriptor.findFieldByName("f_float"), 1.5f); + builder.setField(descriptor.findFieldByName("f_int32"), -123); + builder.setField(descriptor.findFieldByName("f_sint32"), -456); + builder.setField(descriptor.findFieldByName("f_uint32"), 0xFFFFFFFF); + builder.setField(descriptor.findFieldByName("f_fixed32"), 0x80000000); + builder.setField(descriptor.findFieldByName("f_sfixed32"), -654_321); + builder.setField(descriptor.findFieldByName("f_int64"), -9_876_543_210L); + builder.setField(descriptor.findFieldByName("f_sint64"), -123_456_789_012L); + builder.setField(descriptor.findFieldByName("f_uint64"), -1L); + builder.setField(descriptor.findFieldByName("f_fixed64"), Long.MIN_VALUE); + builder.setField(descriptor.findFieldByName("f_sfixed64"), -9_223_372_036_854_775_000L); + builder.setField(descriptor.findFieldByName("f_string"), "string-value"); + builder.setField(descriptor.findFieldByName("f_bytes"), ByteString.copyFromUtf8("bytes-value")); + builder.setField( + descriptor.findFieldByName("f_enum"), + descriptor.getFile().findEnumTypeByName("SampleEnum").findValueByName("SAMPLE_ENUM_SECOND") + ); + + // Build FloatArray for oneof choice + Descriptors.FieldDescriptor floatArrayField = descriptor.findFieldByName("choice_float_array"); + Descriptors.Descriptor floatArrayDescriptor = floatArrayField.getMessageType(); + DynamicMessage.Builder floatArrayBuilder = DynamicMessage.newBuilder(floatArrayDescriptor); + Descriptors.FieldDescriptor floatValuesField = floatArrayDescriptor.findFieldByName("values"); + floatArrayBuilder.addRepeatedField(floatValuesField, 1.1); + floatArrayBuilder.addRepeatedField(floatValuesField, 2.2); + floatArrayBuilder.addRepeatedField(floatValuesField, 3.3); + builder.setField(floatArrayField, floatArrayBuilder.build()); + + Descriptors.FieldDescriptor nestedField = descriptor.findFieldByName("f_message"); + Descriptors.Descriptor nestedDescriptor = nestedField.getMessageType(); + DynamicMessage nestedMessage = DynamicMessage.newBuilder(nestedDescriptor) + .setField(nestedDescriptor.findFieldByName("name"), "nested-name") + .setField(nestedDescriptor.findFieldByName("count"), 7) + .build(); + builder.setField(nestedField, nestedMessage); + + Descriptors.FieldDescriptor repeatedNestedField = descriptor.findFieldByName("f_nested_list"); + builder.addRepeatedField(repeatedNestedField, nestedMessage); + DynamicMessage nestedMessage2 = DynamicMessage.newBuilder(nestedDescriptor) + .setField(nestedDescriptor.findFieldByName("name"), "nested-name-2") + .setField(nestedDescriptor.findFieldByName("count"), 8) + .build(); + builder.addRepeatedField(repeatedNestedField, nestedMessage2); + + Descriptors.FieldDescriptor repeatedField = descriptor.findFieldByName("f_repeated_int32"); + builder.addRepeatedField(repeatedField, 1); + builder.addRepeatedField(repeatedField, 2); + builder.addRepeatedField(repeatedField, 3); + + Descriptors.FieldDescriptor mapField = descriptor.findFieldByName("f_string_int32_map"); + Descriptors.Descriptor entryDescriptor = mapField.getMessageType(); + builder.addRepeatedField(mapField, mapEntry(entryDescriptor, "key1", 11)); + builder.addRepeatedField(mapField, mapEntry(entryDescriptor, "key2", 22)); + + Descriptors.FieldDescriptor nestedMapField = descriptor.findFieldByName("f_string_nested_map"); + Descriptors.Descriptor nestedEntryDescriptor = nestedMapField.getMessageType(); + builder.addRepeatedField(nestedMapField, mapEntry(nestedEntryDescriptor, "nk1", nestedMessage)); + builder.addRepeatedField(nestedMapField, mapEntry(nestedEntryDescriptor, "nk2", nestedMessage2)); + + Timestamp timestamp = Timestamp.newBuilder().setSeconds(1_234_567_890L).setNanos(987_000_000).build(); + builder.setField(descriptor.findFieldByName("f_timestamp"), timestamp); + + Descriptors.FieldDescriptor timestampMapField = descriptor.findFieldByName("f_string_timestamp_map"); + Descriptors.Descriptor timestampEntryDescriptor = timestampMapField.getMessageType(); + Timestamp timestamp1 = Timestamp.newBuilder().setSeconds(1_600_000_000L).setNanos(123_000_000).build(); + Timestamp timestamp2 = Timestamp.newBuilder().setSeconds(1_700_000_000L).setNanos(456_000_000).build(); + builder.addRepeatedField(timestampMapField, mapEntry(timestampEntryDescriptor, "ts1", timestamp1)); + builder.addRepeatedField(timestampMapField, mapEntry(timestampEntryDescriptor, "ts2", timestamp2)); + + return builder.build(); + } + + @Test + void testConvertStandaloneMapField() throws Exception { + String topic = "pb-map-only"; + String subject = topic + "-value"; + + MockSchemaRegistryClient registryClient = new MockSchemaRegistryClient(List.of(new ProtobufSchemaProvider())); + CustomProtobufSchema schema = new CustomProtobufSchema( + "MapOnly", + -1, + null, + null, + MAP_ONLY_PROTO, + List.of(), + Map.of() + ); + int schemaId = registryClient.register(subject, schema); + + ProtoFileElement fileElement = ProtoParser.Companion.parse(ProtoConstants.DEFAULT_LOCATION, MAP_ONLY_PROTO); + DynamicSchema dynamicSchema = ProtobufSchemaParser.toDynamicSchema("MapOnly", fileElement, Collections.emptyMap()); + Descriptors.Descriptor descriptor = dynamicSchema.getMessageDescriptor("MapOnly"); + + DynamicMessage message = buildMapOnlyMessage(descriptor); + ByteBuffer payload = buildConfluentPayload(schemaId, message.toByteArray(), 0); + + ProtobufRegistryConverter converter = new ProtobufRegistryConverter(registryClient, "http://mock:8081", false); + ConversionResult result = converter.convert(topic, payload.asReadOnlyBuffer()); + + GenericRecord record = (GenericRecord) result.getValue(); + List attributeEntries = (List) record.get("attributes"); + Map attributes = attributeEntries.stream() + .map(GenericRecord.class::cast) + .collect(Collectors.toMap( + entry -> entry.get("key").toString(), + entry -> (Integer) entry.get("value") + )); + + assertEquals(Map.of("env", 1, "tier", 2), attributes); + + Schema.Field attributesField = record.getSchema().getField("attributes"); + Schema mapSchema = attributesField.schema(); + assertNotNull(mapSchema.getLogicalType(), "Map field should have logical type"); + assertEquals("map", mapSchema.getLogicalType().getName()); + assertEquals(GenericData.Array.class, record.get("attributes").getClass()); + + org.apache.iceberg.Schema icebergSchema = AvroSchemaUtil.toIceberg(record.getSchema()); + assertEquals(Type.TypeID.MAP, icebergSchema.findField("attributes").type().typeId()); + } + + private static DynamicMessage buildMapOnlyMessage(Descriptors.Descriptor descriptor) { + DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptor); + Descriptors.FieldDescriptor mapField = descriptor.findFieldByName("attributes"); + Descriptors.Descriptor entryDescriptor = mapField.getMessageType(); + builder.addRepeatedField(mapField, mapEntry(entryDescriptor, "env", 1)); + builder.addRepeatedField(mapField, mapEntry(entryDescriptor, "tier", 2)); + return builder.build(); + } + + private static DynamicMessage mapEntry(Descriptors.Descriptor entryDescriptor, Object key, Object value) { + return DynamicMessage.newBuilder(entryDescriptor) + .setField(entryDescriptor.findFieldByName("key"), key) + .setField(entryDescriptor.findFieldByName("value"), value) + .build(); + } + + private static ByteBuffer buildConfluentPayload(int schemaId, byte[] messageBytes, int... messageIndexes) { + byte[] indexBytes = encodeMessageIndexes(messageIndexes); + ByteBuffer buffer = ByteBuffer.allocate(1 + Integer.BYTES + indexBytes.length + messageBytes.length); + buffer.put((byte) 0); + buffer.putInt(schemaId); + buffer.put(indexBytes); + buffer.put(messageBytes); + buffer.flip(); + return buffer; + } + + private static byte[] encodeMessageIndexes(int... indexes) { + if (indexes == null || indexes.length == 0) { + return new byte[]{0}; + } + ByteBuffer buffer = ByteBuffer.allocate(5 * (indexes.length + 1)); + ByteUtils.writeVarint(indexes.length, buffer); + for (int index : indexes) { + ByteUtils.writeVarint(index, buffer); + } + buffer.flip(); + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + return bytes; + } + + private static void assertPrimitiveFields(GenericRecord record) { + assertEquals(true, record.get("f_bool")); + assertEquals(123.456d, (double) record.get("f_double"), 1e-6); + assertEquals(1.5f, (Float) record.get("f_float"), 1e-6); + assertEquals(-123, ((Integer) record.get("f_int32")).intValue()); + assertEquals(-456, ((Integer) record.get("f_sint32")).intValue()); + int unsigned32 = ((Integer) record.get("f_uint32")).intValue(); + assertEquals("4294967295", Long.toString(Integer.toUnsignedLong(unsigned32)), "f_uint32 preserves unsigned semantics despite signed storage"); + assertEquals(Integer.MIN_VALUE, ((Integer) record.get("f_fixed32")).intValue()); + assertEquals(-654_321, ((Integer) record.get("f_sfixed32")).intValue()); + assertEquals(-9_876_543_210L, ((Long) record.get("f_int64")).longValue()); + assertEquals(-123_456_789_012L, ((Long) record.get("f_sint64")).longValue()); + long uint64 = ((Long) record.get("f_uint64")).longValue(); + assertEquals("18446744073709551615", Long.toUnsignedString(uint64), "f_uint64 preserves unsigned semantics despite signed storage"); + assertEquals(Long.MIN_VALUE, ((Long) record.get("f_fixed64")).longValue()); + assertEquals(-9_223_372_036_854_775_000L, ((Long) record.get("f_sfixed64")).longValue()); + assertEquals("string-value", record.get("f_string").toString()); + + ByteBuffer bytesBuffer = ((ByteBuffer) record.get("f_bytes")).duplicate(); + byte[] bytes = new byte[bytesBuffer.remaining()]; + bytesBuffer.get(bytes); + assertEquals("bytes-value", new String(bytes, StandardCharsets.UTF_8)); + + assertEquals("SAMPLE_ENUM_SECOND", record.get("f_enum").toString()); + } + + private static void assertRepeatedAndMapFields(GenericRecord record) { + List repeated = ((List) record.get("f_repeated_int32")).stream() + .map(value -> (Integer) value) + .collect(Collectors.toList()); + assertEquals(List.of(1, 2, 3), repeated); + + List mapEntries = (List) record.get("f_string_int32_map"); + Map map = mapEntries.stream() + .map(GenericRecord.class::cast) + .collect(Collectors.toMap( + entry -> entry.get("key").toString(), + entry -> (Integer) entry.get("value") + )); + assertEquals(Map.of("key1", 11, "key2", 22), map); + + List nestedList = (List) getField(record, "f_nested_list", "fNestedList"); + List nestedNames = nestedList.stream() + .map(GenericRecord.class::cast) + .map(n -> n.get("name").toString()) + .collect(Collectors.toList()); + assertEquals(List.of("nested-name", "nested-name-2"), nestedNames); + + List nestedMapEntries = (List) getField(record, "f_string_nested_map", "fStringNestedMap"); + Map nestedMap = nestedMapEntries.stream() + .map(GenericRecord.class::cast) + .collect(Collectors.toMap( + entry -> entry.get("key").toString(), + entry -> ((GenericRecord) entry.get("value")).get("name").toString() + )); + assertEquals(Map.of("nk1", "nested-name", "nk2", "nested-name-2"), nestedMap); + + List timestampMapEntries = (List) getField(record, "f_string_timestamp_map", "fStringTimestampMap"); + Map timestampMap = timestampMapEntries.stream() + .map(GenericRecord.class::cast) + .collect(Collectors.toMap( + entry -> entry.get("key").toString(), + entry -> (Long) entry.get("value") + )); + long expectedMicros1 = 1_600_000_000_000_000L + 123_000; + long expectedMicros2 = 1_700_000_000_000_000L + 456_000; + assertEquals(Map.of("ts1", expectedMicros1, "ts2", expectedMicros2), timestampMap); + } + + private static void assertNestedAndTimestamp(GenericRecord record) { + GenericRecord nested = (GenericRecord) getField(record, "f_message", "fMessage"); + assertEquals("nested-name", nested.get("name").toString()); + assertEquals(7, nested.get("count")); + + long expectedMicros = 1_234_567_890_000_000L + 987_000; + assertEquals(expectedMicros, ((Long) record.get("f_timestamp")).longValue()); + + // Optional field should fall back to proto3 default (empty string) + assertEquals("", getField(record, "f_optional_string", "fOptionalString").toString()); + + // Verify oneof with complex FloatArray type + GenericRecord floatArrayValue = (GenericRecord) getField(record, "choice_float_array", "floatArray"); + List floatValues = (List) floatArrayValue.get("values"); + List expectedFloats = List.of(1.1, 2.2, 3.3); + assertEquals(expectedFloats.size(), floatValues.size()); + for (int i = 0; i < expectedFloats.size(); i++) { + assertEquals(expectedFloats.get(i), (Double) floatValues.get(i), 1e-6); + } + } + + private static Object getField(GenericRecord record, String... candidateNames) { + for (String name : candidateNames) { + if (record.getSchema().getField(name) != null) { + return record.get(name); + } + } + throw new IllegalArgumentException("Field not found in schema: " + String.join(", ", candidateNames)); + } +} diff --git a/core/src/test/java/kafka/automq/table/process/convert/ProtobufRegistryConverterUnitTest.java b/core/src/test/java/kafka/automq/table/process/convert/ProtobufRegistryConverterUnitTest.java new file mode 100644 index 0000000000..a6db377ea3 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/convert/ProtobufRegistryConverterUnitTest.java @@ -0,0 +1,719 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package kafka.automq.table.process.convert; + +import kafka.automq.table.binder.RecordBinder; +import kafka.automq.table.deserializer.proto.CustomProtobufSchema; +import kafka.automq.table.deserializer.proto.ProtobufSchemaProvider; +import kafka.automq.table.deserializer.proto.parse.ProtobufSchemaParser; +import kafka.automq.table.deserializer.proto.parse.converter.ProtoConstants; +import kafka.automq.table.deserializer.proto.schema.DynamicSchema; +import kafka.automq.table.process.ConversionResult; + +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import com.google.protobuf.DynamicMessage; +import com.google.protobuf.Timestamp; +import com.squareup.wire.schema.internal.parser.ProtoFileElement; +import com.squareup.wire.schema.internal.parser.ProtoParser; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; +import org.apache.iceberg.Table; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.io.TaskWriter; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient; + +import static kafka.automq.table.binder.AvroRecordBinderTypeTest.createTableWriter; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; + +@Tag("S3Unit") +public class ProtobufRegistryConverterUnitTest { + + private static final String BASIC_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + message BasicRecord { + bool active = 1; + int32 score = 2; + uint32 quota = 3; + int64 total = 4; + uint64 big_total = 5; + float ratio = 6; + double precise = 7; + string name = 8; + bytes payload = 9; + Status status = 10; + Nested meta = 11; + } + + message Nested { + string note = 1; + } + + enum Status { + STATUS_UNSPECIFIED = 0; + STATUS_READY = 1; + } + """; + + private static final String COLLECTION_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + message CollectionRecord { + repeated string tags = 1; + repeated Item notes = 2; + repeated Wrap wrappers = 3; + map counters = 4; + map keyed_items = 5; + map wrap_map = 6; + } + + message Item { + string value = 1; + } + + message Wrap { + Item item = 1; + repeated Item items = 2; + } + """; + + private static final String OPTIONAL_COLLECTION_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + import \"google/protobuf/timestamp.proto\"; + + message OptionalCollectionRecord { + optional Wrapper opt_wrapper = 1; + optional IntStringMap opt_int_map = 2; + optional Item opt_item = 3; + optional WrapMapHolder opt_wrap_map = 4; + optional google.protobuf.Timestamp opt_ts = 5; + } + + message Item { + string value = 1; + } + + message Wrapper { + repeated Item items = 1; + } + + message WrapMapHolder { + map entries = 1; + } + + message IntStringMap { + map entries = 1; + } + """; + + private static final String ADVANCED_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + import \"google/protobuf/timestamp.proto\"; + + message AdvancedRecord { + optional string opt_str = 1; + optional int32 opt_int = 2; + optional Ref opt_ref = 3; + + oneof selection { + string selection_str = 4; + int32 selection_int = 5; + Ref selection_ref = 6; + Bag selection_bag = 7; + MapHolder selection_map = 8; + IntMapHolder selection_int_map = 9; + } + + google.protobuf.Timestamp event_time = 10; + Ref direct = 11; + repeated Ref refs = 12; + } + + message Ref { + string name = 1; + } + + message Bag { + repeated Ref refs = 1; + } + + message MapHolder { + map entries = 1; + } + + message IntMapHolder { + map entries = 1; + } + """; + + private static final String RECURSIVE_PROTO = """ + syntax = \"proto3\"; + + package kafka.automq.table.process.proto; + + message Node { + string id = 1; + Child child = 2; + } + + message Child { + Node leaf = 1; + } + """; + + @Test + void convertBasicTypesRecord() throws Exception { + String topic = "proto-basic"; + ConversionResult result = convert(topic, BASIC_PROTO, "BasicRecord", builder -> { + builder.setField(builder.getDescriptorForType().findFieldByName("active"), true); + builder.setField(builder.getDescriptorForType().findFieldByName("score"), -10); + builder.setField(builder.getDescriptorForType().findFieldByName("quota"), -1); // uint32 max + builder.setField(builder.getDescriptorForType().findFieldByName("total"), -123456789L); + builder.setField(builder.getDescriptorForType().findFieldByName("big_total"), -1L); + builder.setField(builder.getDescriptorForType().findFieldByName("ratio"), 1.5f); + builder.setField(builder.getDescriptorForType().findFieldByName("precise"), 3.14159d); + builder.setField(builder.getDescriptorForType().findFieldByName("name"), "basic-name"); + builder.setField( + builder.getDescriptorForType().findFieldByName("payload"), + ByteString.copyFromUtf8("payload-bytes") + ); + builder.setField( + builder.getDescriptorForType().findFieldByName("status"), + builder.getDescriptorForType().getFile().findEnumTypeByName("Status").findValueByName("STATUS_READY") + ); + Descriptors.FieldDescriptor nestedField = builder.getDescriptorForType().findFieldByName("meta"); + builder.setField(nestedField, nestedMessage(nestedField.getMessageType(), "note-value")); + }); + + GenericRecord record = (GenericRecord) result.getValue(); + assertEquals(true, record.get("active")); + assertEquals(-10, record.get("score")); + int quotaSigned = (Integer) record.get("quota"); + assertEquals("4294967295", Long.toUnsignedString(Integer.toUnsignedLong(quotaSigned))); + assertEquals(-123456789L, record.get("total")); + long bigTotal = (Long) record.get("big_total"); + assertEquals("18446744073709551615", Long.toUnsignedString(bigTotal)); + assertEquals(1.5f, (Float) record.get("ratio"), 1e-6); + assertEquals(3.14159d, (Double) record.get("precise"), 1e-9); + assertEquals("basic-name", record.get("name").toString()); + assertEquals("payload-bytes", utf8(record.get("payload"))); + assertEquals("STATUS_READY", record.get("status").toString()); + assertEquals("note-value", ((GenericRecord) record.get("meta")).get("note").toString()); + + bindAndWrite(record); + } + + @Test + void convertCollectionsRecord() throws Exception { + String topic = "proto-collections"; + ConversionResult result = convert(topic, COLLECTION_PROTO, "CollectionRecord", builder -> { + Descriptors.FieldDescriptor tagsFd = builder.getDescriptorForType().findFieldByName("tags"); + builder.addRepeatedField(tagsFd, "alpha"); + builder.addRepeatedField(tagsFd, "beta"); + + Descriptors.FieldDescriptor notesFd = builder.getDescriptorForType().findFieldByName("notes"); + Descriptors.Descriptor itemDesc = notesFd.getMessageType(); + builder.addRepeatedField(notesFd, nestedMessage(itemDesc, "note-1")); + builder.addRepeatedField(notesFd, nestedMessage(itemDesc, "note-2")); + + Descriptors.FieldDescriptor wrappersFd = builder.getDescriptorForType().findFieldByName("wrappers"); + Descriptors.Descriptor wrapDesc = wrappersFd.getMessageType(); + Descriptors.Descriptor wrapItemDesc = wrapDesc.findFieldByName("item").getMessageType(); + builder.addRepeatedField(wrappersFd, wrapMessage(wrapDesc, wrapItemDesc, "w1-a", List.of("w1-b"))); + builder.addRepeatedField(wrappersFd, wrapMessage(wrapDesc, wrapItemDesc, "w2-a", List.of("w2-b", "w2-c"))); + + Descriptors.FieldDescriptor countersFd = builder.getDescriptorForType().findFieldByName("counters"); + Descriptors.Descriptor countersEntry = countersFd.getMessageType(); + builder.addRepeatedField(countersFd, mapEntry(countersEntry, "k1", 10)); + builder.addRepeatedField(countersFd, mapEntry(countersEntry, "k2", 20)); + + Descriptors.FieldDescriptor keyedFd = builder.getDescriptorForType().findFieldByName("keyed_items"); + Descriptors.Descriptor keyedEntry = keyedFd.getMessageType(); + builder.addRepeatedField(keyedFd, mapEntry(keyedEntry, 1, nestedMessage(keyedEntry.findFieldByName("value").getMessageType(), "v1"))); + builder.addRepeatedField(keyedFd, mapEntry(keyedEntry, 2, nestedMessage(keyedEntry.findFieldByName("value").getMessageType(), "v2"))); + + Descriptors.FieldDescriptor wrapMapFd = builder.getDescriptorForType().findFieldByName("wrap_map"); + Descriptors.Descriptor wrapMapEntry = wrapMapFd.getMessageType(); + Descriptors.Descriptor wrapMapValueDesc = wrapMapEntry.findFieldByName("value").getMessageType(); + Descriptors.Descriptor wrapMapItemDesc = wrapMapValueDesc.findFieldByName("item").getMessageType(); + builder.addRepeatedField(wrapMapFd, mapEntry( + wrapMapEntry, + "wm1", + wrapMessage(wrapMapValueDesc, wrapMapItemDesc, "wm1-a", List.of("wm1-b")) + )); + builder.addRepeatedField(wrapMapFd, mapEntry( + wrapMapEntry, + "wm2", + wrapMessage(wrapMapValueDesc, wrapMapItemDesc, "wm2-a", List.of("wm2-b", "wm2-c")) + )); + }); + + GenericRecord record = (GenericRecord) result.getValue(); + List tags = (List) record.get("tags"); + assertEquals(List.of("alpha", "beta"), tags.stream().map(Object::toString).collect(Collectors.toList())); + + List notes = (List) record.get("notes"); + assertEquals(List.of("note-1", "note-2"), notes.stream() + .map(GenericRecord.class::cast) + .map(r -> r.get("value").toString()) + .collect(Collectors.toList())); + + Map counters = logicalMapToMap(record.get("counters")); + assertEquals(Map.of("k1", 10, "k2", 20), counters); + + Map keyed = logicalMapToMap(record.get("keyed_items")); + assertEquals(Map.of(1, "v1", 2, "v2"), keyed); + + List wrappers = (List) record.get("wrappers"); + assertEquals(List.of("w1-a", "w2-a"), wrappers.stream() + .map(GenericRecord.class::cast) + .map(r -> ((GenericRecord) r.get("item")).get("value").toString()) + .collect(Collectors.toList())); + assertEquals(List.of( + List.of("w1-b"), + List.of("w2-b", "w2-c") + ), wrappers.stream() + .map(GenericRecord.class::cast) + .map(r -> (List) r.get("items")) + .map(lst -> lst.stream() + .map(GenericRecord.class::cast) + .map(it -> it.get("value").toString()) + .collect(Collectors.toList())) + .collect(Collectors.toList())); + + Map wrapMap = logicalMapToMap(record.get("wrap_map")); + assertEquals("wm1-a", ((GenericRecord) ((GenericRecord) wrapMap.get("wm1")).get("item")).get("value").toString()); + assertEquals("wm2-a", ((GenericRecord) ((GenericRecord) wrapMap.get("wm2")).get("item")).get("value").toString()); + + bindAndWrite(record); + } + + @Test + void convertAdvancedRecord() throws Exception { + String topic = "proto-advanced"; + ConversionResult result = convert(topic, ADVANCED_PROTO, "AdvancedRecord", builder -> { + builder.setField(builder.getDescriptorForType().findFieldByName("opt_str"), "optional-value"); + builder.setField(builder.getDescriptorForType().findFieldByName("opt_int"), 99); + Descriptors.FieldDescriptor optRefFd = builder.getDescriptorForType().findFieldByName("opt_ref"); + builder.setField(optRefFd, nestedMessage(optRefFd.getMessageType(), "opt-ref")); + + // choose oneof map branch via MapHolder; other branches should remain null + Descriptors.FieldDescriptor selMapFd = builder.getDescriptorForType().findFieldByName("selection_map"); + Descriptors.Descriptor mapHolderDesc = selMapFd.getMessageType(); + Descriptors.FieldDescriptor entriesFd = mapHolderDesc.findFieldByName("entries"); + Descriptors.Descriptor entryDesc = entriesFd.getMessageType(); + DynamicMessage.Builder holderBuilder = DynamicMessage.newBuilder(mapHolderDesc); + holderBuilder.addRepeatedField(entriesFd, mapEntry(entryDesc, "a", 1)); + holderBuilder.addRepeatedField(entriesFd, mapEntry(entryDesc, "b", 2)); + builder.setField(selMapFd, holderBuilder.build()); + Timestamp timestamp = Timestamp.newBuilder().setSeconds(1234L).setNanos(567000000).build(); + builder.setField(builder.getDescriptorForType().findFieldByName("event_time"), timestamp); + + Descriptors.FieldDescriptor refField = builder.getDescriptorForType().findFieldByName("direct"); + builder.setField(refField, nestedMessage(refField.getMessageType(), "parent")); + + Descriptors.FieldDescriptor refsField = builder.getDescriptorForType().findFieldByName("refs"); + Descriptors.Descriptor refDescriptor = refsField.getMessageType(); + builder.addRepeatedField(refsField, nestedMessage(refDescriptor, "child-1")); + builder.addRepeatedField(refsField, nestedMessage(refDescriptor, "child-2")); + }); + + GenericRecord record = (GenericRecord) result.getValue(); + Schema optionalSchema = record.getSchema().getField("opt_str").schema(); + assertEquals(Schema.Type.UNION, optionalSchema.getType()); + assertEquals(Schema.Type.STRING, optionalSchema.getTypes().get(0).getType()); + assertEquals("optional-value", record.get("opt_str").toString()); + assertEquals(99, record.get("opt_int")); + assertEquals("opt-ref", ((GenericRecord) record.get("opt_ref")).get("name").toString()); + + GenericRecord selMapRecord = (GenericRecord) record.get("selection_map"); + Map selMap = logicalMapToMap(selMapRecord.get("entries")); + assertEquals(Map.of("a", 1, "b", 2), selMap); + assertEquals(null, record.get("selection_ref")); + assertEquals(null, record.get("selection_str")); + assertEquals(null, record.get("selection_int")); + assertEquals(null, record.get("selection_bag")); + + long expectedMicros = 1234_000_000L + 567_000; + assertEquals(expectedMicros, record.get("event_time")); + assertEquals("parent", ((GenericRecord) record.get("direct")).get("name").toString()); + + List refs = (List) record.get("refs"); + assertEquals(List.of("child-1", "child-2"), refs.stream() + .map(GenericRecord.class::cast) + .map(r -> r.get("name").toString()) + .collect(Collectors.toList())); + + bindAndWrite(record); + } + + @Test + void convertAdvancedOneofStringIntRefBag() throws Exception { + // string branch + ConversionResult stringResult = convert("proto-adv-oneof-str", ADVANCED_PROTO, "AdvancedRecord", b -> + b.setField(b.getDescriptorForType().findFieldByName("selection_str"), "sel-str")); + GenericRecord stringRec = (GenericRecord) stringResult.getValue(); + assertEquals("sel-str", stringRec.get("selection_str")); + assertEquals(null, stringRec.get("selection_int")); + assertEquals(null, stringRec.get("selection_ref")); + assertEquals(null, stringRec.get("selection_bag")); + bindAndWrite((GenericRecord) stringResult.getValue()); + + // int branch + ConversionResult intResult = convert("proto-adv-oneof-int", ADVANCED_PROTO, "AdvancedRecord", b -> + b.setField(b.getDescriptorForType().findFieldByName("selection_int"), 123)); + GenericRecord intRec = (GenericRecord) intResult.getValue(); + assertEquals(123, intRec.get("selection_int")); + assertEquals(null, intRec.get("selection_str")); + assertEquals(null, intRec.get("selection_ref")); + assertEquals(null, intRec.get("selection_bag")); + bindAndWrite((GenericRecord) intResult.getValue()); + + // ref branch + ConversionResult refResult = convert("proto-adv-oneof-ref", ADVANCED_PROTO, "AdvancedRecord", b -> { + Descriptors.FieldDescriptor fd = b.getDescriptorForType().findFieldByName("selection_ref"); + b.setField(fd, nestedMessage(fd.getMessageType(), "sel-ref")); + }); + GenericRecord refRec = (GenericRecord) refResult.getValue(); + assertEquals("sel-ref", ((GenericRecord) refRec.get("selection_ref")).get("name").toString()); + assertEquals(null, refRec.get("selection_str")); + assertEquals(null, refRec.get("selection_int")); + assertEquals(null, refRec.get("selection_bag")); + bindAndWrite((GenericRecord) refResult.getValue()); + + // bag branch (contains repeated refs) + ConversionResult bagResult = convert("proto-adv-oneof-bag", ADVANCED_PROTO, "AdvancedRecord", b -> { + Descriptors.FieldDescriptor fd = b.getDescriptorForType().findFieldByName("selection_bag"); + Descriptors.Descriptor bagDesc = fd.getMessageType(); + Descriptors.FieldDescriptor refsFd = bagDesc.findFieldByName("refs"); + DynamicMessage.Builder bagBuilder = DynamicMessage.newBuilder(bagDesc); + bagBuilder.addRepeatedField(refsFd, nestedMessage(refsFd.getMessageType(), "b1")); + bagBuilder.addRepeatedField(refsFd, nestedMessage(refsFd.getMessageType(), "b2")); + b.setField(fd, bagBuilder.build()); + }); + GenericRecord bagRec = (GenericRecord) bagResult.getValue(); + List bagRefs = (List) ((GenericRecord) bagRec.get("selection_bag")).get("refs"); + assertEquals(List.of("b1", "b2"), bagRefs.stream() + .map(GenericRecord.class::cast) + .map(r -> r.get("name").toString()) + .collect(Collectors.toList())); + assertEquals(null, bagRec.get("selection_str")); + assertEquals(null, bagRec.get("selection_int")); + assertEquals(null, bagRec.get("selection_ref")); + bindAndWrite((GenericRecord) bagResult.getValue()); + + // int map branch (map) + ConversionResult intMapResult = convert("proto-adv-oneof-intmap", ADVANCED_PROTO, "AdvancedRecord", b -> { + Descriptors.FieldDescriptor fd = b.getDescriptorForType().findFieldByName("selection_int_map"); + Descriptors.Descriptor holderDesc = fd.getMessageType(); + Descriptors.FieldDescriptor entriesFd = holderDesc.findFieldByName("entries"); + Descriptors.Descriptor entryDesc = entriesFd.getMessageType(); + DynamicMessage.Builder holder = DynamicMessage.newBuilder(holderDesc); + holder.addRepeatedField(entriesFd, mapEntry(entryDesc, 10, "x")); + holder.addRepeatedField(entriesFd, mapEntry(entryDesc, 20, "y")); + b.setField(fd, holder.build()); + }); + GenericRecord intMapRec = (GenericRecord) intMapResult.getValue(); + Map intMaps = logicalMapToMap(((GenericRecord) intMapRec.get("selection_int_map")).get("entries")); + assertEquals(Map.of(10, "x", 20, "y"), intMaps); + assertEquals(null, intMapRec.get("selection_str")); + assertEquals(null, intMapRec.get("selection_int")); + assertEquals(null, intMapRec.get("selection_ref")); + assertEquals(null, intMapRec.get("selection_bag")); + bindAndWrite((GenericRecord) intMapResult.getValue()); + } + + @Test + void convertOptionalCollectionsRecord() throws Exception { + String topic = "proto-optional-collections"; + ConversionResult result = convert(topic, OPTIONAL_COLLECTION_PROTO, "OptionalCollectionRecord", builder -> { + Descriptors.FieldDescriptor wrapperFd = builder.getDescriptorForType().findFieldByName("opt_wrapper"); + Descriptors.Descriptor wrapperDesc = wrapperFd.getMessageType(); + Descriptors.Descriptor itemDesc = wrapperDesc.findFieldByName("items").getMessageType(); + DynamicMessage.Builder wrapperBuilder = DynamicMessage.newBuilder(wrapperDesc); + wrapperBuilder.addRepeatedField(wrapperDesc.findFieldByName("items"), nestedMessage(itemDesc, "i1")); + wrapperBuilder.addRepeatedField(wrapperDesc.findFieldByName("items"), nestedMessage(itemDesc, "i2")); + builder.setField(wrapperFd, wrapperBuilder.build()); + + // leave opt_int_map unset to validate optional-map -> null + + Descriptors.FieldDescriptor optItemFd = builder.getDescriptorForType().findFieldByName("opt_item"); + builder.setField(optItemFd, nestedMessage(optItemFd.getMessageType(), "single")); + }); + + GenericRecord record = (GenericRecord) result.getValue(); + // opt_wrapper union present + assertEquals(Schema.Type.UNION, record.getSchema().getField("opt_wrapper").schema().getType()); + List items = (List) ((GenericRecord) record.get("opt_wrapper")).get("items"); + assertEquals(List.of("i1", "i2"), items.stream().map(GenericRecord.class::cast).map(r -> r.get("value").toString()).collect(Collectors.toList())); + + assertEquals(null, record.get("opt_int_map")); + + GenericRecord optItem = (GenericRecord) record.get("opt_item"); + assertEquals("single", optItem.get("value").toString()); + } + + @Test + void convertOptionalCollectionsRecordWithMap() throws Exception { + String topic = "proto-optional-collections-map"; + ConversionResult result = convert(topic, OPTIONAL_COLLECTION_PROTO, "OptionalCollectionRecord", builder -> { + Descriptors.FieldDescriptor optMapFd = builder.getDescriptorForType().findFieldByName("opt_int_map"); + Descriptors.Descriptor holderDesc = optMapFd.getMessageType(); + Descriptors.FieldDescriptor entriesFd = holderDesc.findFieldByName("entries"); + Descriptors.Descriptor entryDesc = entriesFd.getMessageType(); + DynamicMessage.Builder holder = DynamicMessage.newBuilder(holderDesc); + holder.addRepeatedField(entriesFd, mapEntry(entryDesc, 7, "v7")); + holder.addRepeatedField(entriesFd, mapEntry(entryDesc, 8, "v8")); + builder.setField(optMapFd, holder.build()); + + Descriptors.FieldDescriptor wrapMapFd = builder.getDescriptorForType().findFieldByName("opt_wrap_map"); + Descriptors.Descriptor wrapMapDesc = wrapMapFd.getMessageType(); + Descriptors.FieldDescriptor wrapEntriesFd = wrapMapDesc.findFieldByName("entries"); + Descriptors.Descriptor wrapEntryDesc = wrapEntriesFd.getMessageType(); + Descriptors.Descriptor wrapValueDesc = wrapEntryDesc.findFieldByName("value").getMessageType(); + Descriptors.Descriptor wrapItemDesc = wrapValueDesc.findFieldByName("items").getMessageType(); + + DynamicMessage.Builder wrapHolder = DynamicMessage.newBuilder(wrapMapDesc); + wrapHolder.addRepeatedField(wrapEntriesFd, mapEntry( + wrapEntryDesc, + "wkey1", + wrapMessage(wrapValueDesc, wrapItemDesc, "wm1-a", List.of("wm1-b")) + )); + builder.setField(wrapMapFd, wrapHolder.build()); + + // optional timestamp + Timestamp ts = Timestamp.newBuilder().setSeconds(10L).setNanos(500_000_000).build(); + builder.setField(builder.getDescriptorForType().findFieldByName("opt_ts"), ts); + }); + + GenericRecord record = (GenericRecord) result.getValue(); + GenericRecord optIntMap = (GenericRecord) record.get("opt_int_map"); + Map map = logicalMapToMap(optIntMap.get("entries")); + assertEquals(Map.of(7, "v7", 8, "v8"), map); + + Schema.Field optWrapField = record.getSchema().getField("opt_wrap_map"); + assertEquals(Schema.Type.UNION, optWrapField.schema().getType()); + GenericRecord optWrapMap = (GenericRecord) record.get("opt_wrap_map"); + Map wrapEntries = logicalMapToMap(optWrapMap.get("entries")); + GenericRecord wrapper = wrapEntries.get("wkey1"); + List wrapItems = (List) wrapper.get("items"); + assertEquals(List.of("wm1-b"), wrapItems.stream() + .map(GenericRecord.class::cast) + .map(item -> item.get("value").toString()) + .collect(Collectors.toList())); + + assertEquals(10_500_000L, record.get("opt_ts")); + + + } + + @Test + void convertRecursiveRecord() throws Exception { + String topic = "proto-recursive"; + ConversionResult result = convert(topic, RECURSIVE_PROTO, "Node", builder -> { + builder.setField(builder.getDescriptorForType().findFieldByName("id"), "root"); + Descriptors.FieldDescriptor childFd = builder.getDescriptorForType().findFieldByName("child"); + Descriptors.Descriptor childDesc = childFd.getMessageType(); + Descriptors.FieldDescriptor leafFd = childDesc.findFieldByName("leaf"); + Descriptors.Descriptor nodeDesc = leafFd.getMessageType(); + + DynamicMessage leaf = DynamicMessage.newBuilder(nodeDesc) + .setField(nodeDesc.findFieldByName("id"), "leaf") + .build(); + DynamicMessage child = DynamicMessage.newBuilder(childDesc) + .setField(leafFd, leaf) + .build(); + builder.setField(childFd, child); + }); + + GenericRecord record = (GenericRecord) result.getValue(); + assertEquals("root", record.get("id").toString()); + GenericRecord child = (GenericRecord) record.get("child"); + GenericRecord leaf = (GenericRecord) child.get("leaf"); + assertEquals("leaf", leaf.get("id").toString()); + + assertThrows(IllegalStateException.class, () -> bindAndWrite(record)); + } + + private ConversionResult convert(String topic, String proto, String messageName, Consumer messageConfigurer) throws Exception { + MockSchemaRegistryClient registryClient = new MockSchemaRegistryClient(List.of(new ProtobufSchemaProvider())); + CustomProtobufSchema schema = new CustomProtobufSchema( + messageName, + -1, + null, + null, + proto, + List.of(), + Map.of() + ); + int schemaId = registryClient.register(topic + "-value", schema); + + ProtoFileElement fileElement = ProtoParser.Companion.parse(ProtoConstants.DEFAULT_LOCATION, proto); + DynamicSchema dynamicSchema = ProtobufSchemaParser.toDynamicSchema(messageName, fileElement, Collections.emptyMap()); + Descriptors.Descriptor descriptor = dynamicSchema.getMessageDescriptor(messageName); + + DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptor); + messageConfigurer.accept(builder); + DynamicMessage message = builder.build(); + + ByteBuffer payload = buildConfluentPayload(schemaId, message.toByteArray(), 0); + ProtobufRegistryConverter converter = new ProtobufRegistryConverter(registryClient, "http://mock:8081", false); + ConversionResult result = converter.convert(topic, payload.asReadOnlyBuffer()); + + ConversionResult cached = converter.convert(topic, payload.asReadOnlyBuffer()); + assertSame(result.getSchema(), cached.getSchema()); + return result; + } + + private void bindAndWrite(GenericRecord record) { + org.apache.iceberg.Schema iceberg = AvroSchemaUtil.toIceberg(record.getSchema()); + RecordBinder binder = new RecordBinder(iceberg, record.getSchema()); + Record icebergRecord = binder.bind(record); + assertDoesNotThrow(() -> testSendRecord(iceberg, icebergRecord)); + } + + private static String utf8(Object value) { + ByteBuffer buffer = ((ByteBuffer) value).duplicate(); + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + return new String(bytes, StandardCharsets.UTF_8); + } + + private static DynamicMessage nestedMessage(Descriptors.Descriptor descriptor, String value) { + return DynamicMessage.newBuilder(descriptor) + .setField(descriptor.findFieldByName("note") != null ? descriptor.findFieldByName("note") : descriptor.findFieldByName("value") != null + ? descriptor.findFieldByName("value") : descriptor.findFieldByName("name"), value) + .build(); + } + + private static DynamicMessage wrapMessage(Descriptors.Descriptor wrapDesc, Descriptors.Descriptor itemDesc, String itemValue, List itemListValues) { + DynamicMessage.Builder wrapBuilder = DynamicMessage.newBuilder(wrapDesc); + Descriptors.FieldDescriptor itemField = wrapDesc.findFieldByName("item"); + if (itemField != null) { + wrapBuilder.setField(itemField, nestedMessage(itemDesc, itemValue)); + } + Descriptors.FieldDescriptor itemsFd = wrapDesc.findFieldByName("items"); + if (itemsFd != null) { + for (String v : itemListValues) { + wrapBuilder.addRepeatedField(itemsFd, nestedMessage(itemDesc, v)); + } + } + return wrapBuilder.build(); + } + + private static DynamicMessage mapEntry(Descriptors.Descriptor descriptor, Object key, Object value) { + DynamicMessage.Builder builder = DynamicMessage.newBuilder(descriptor); + builder.setField(descriptor.findFieldByName("key"), key); + builder.setField(descriptor.findFieldByName("value"), value); + return builder.build(); + } + + private static Map logicalMapToMap(Object logicalMap) { + List entries = (List) logicalMap; + return entries.stream() + .map(GenericRecord.class::cast) + .collect(Collectors.toMap( + entry -> (K) entry.get("key"), + entry -> { + Object value = entry.get("value"); + if (value instanceof GenericRecord) { + GenericRecord record = (GenericRecord) value; + if (record.getSchema().getField("value") != null) { + return (V) record.get("value").toString(); + } + if (record.getSchema().getField("name") != null) { + return (V) record.get("name").toString(); + } + } + return (V) value; + } + )); + } + + private static ByteBuffer buildConfluentPayload(int schemaId, byte[] messageBytes, int... messageIndexes) { + byte[] indexBytes = encodeMessageIndexes(messageIndexes); + ByteBuffer buffer = ByteBuffer.allocate(1 + Integer.BYTES + indexBytes.length + messageBytes.length); + buffer.put((byte) 0); + buffer.putInt(schemaId); + buffer.put(indexBytes); + buffer.put(messageBytes); + buffer.flip(); + return buffer; + } + + private static byte[] encodeMessageIndexes(int... indexes) { + if (indexes == null || indexes.length == 0) { + return new byte[]{0}; + } + ByteBuffer buffer = ByteBuffer.allocate(5 * (indexes.length + 1)); + org.apache.kafka.common.utils.ByteUtils.writeVarint(indexes.length, buffer); + for (int index : indexes) { + org.apache.kafka.common.utils.ByteUtils.writeVarint(index, buffer); + } + buffer.flip(); + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + return bytes; + } + + private void testSendRecord(org.apache.iceberg.Schema schema, Record record) { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("test", ImmutableMap.of()); + catalog.createNamespace(Namespace.of("default")); + Table table = catalog.createTable(TableIdentifier.of(Namespace.of("default"), "scenario"), schema); + TaskWriter writer = createTableWriter(table); + try { + writer.write(record); + } catch (IOException e) { + throw new RuntimeException(e); + } + } +} diff --git a/core/src/test/java/kafka/automq/table/process/proto/PersonProto.java b/core/src/test/java/kafka/automq/table/process/proto/PersonProto.java new file mode 100644 index 0000000000..b116ae4399 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/proto/PersonProto.java @@ -0,0 +1,2474 @@ +// Generated by the protocol buffer compiler. DO NOT EDIT! +// source: proto/person.proto + +package kafka.automq.table.process.proto; + +public final class PersonProto { + private PersonProto() {} + public static void registerAllExtensions( + com.google.protobuf.ExtensionRegistryLite registry) { + } + + public static void registerAllExtensions( + com.google.protobuf.ExtensionRegistry registry) { + registerAllExtensions( + (com.google.protobuf.ExtensionRegistryLite) registry); + } + public interface AddressOrBuilder extends + // @@protoc_insertion_point(interface_extends:kafka.automq.table.process.proto.Address) + com.google.protobuf.MessageOrBuilder { + + /** + * string street = 1; + * @return The street. + */ + java.lang.String getStreet(); + /** + * string street = 1; + * @return The bytes for street. + */ + com.google.protobuf.ByteString + getStreetBytes(); + + /** + * string city = 2; + * @return The city. + */ + java.lang.String getCity(); + /** + * string city = 2; + * @return The bytes for city. + */ + com.google.protobuf.ByteString + getCityBytes(); + } + /** + * Protobuf type {@code kafka.automq.table.process.proto.Address} + */ + public static final class Address extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:kafka.automq.table.process.proto.Address) + AddressOrBuilder { + private static final long serialVersionUID = 0L; + // Use Address.newBuilder() to construct. + private Address(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private Address() { + street_ = ""; + city_ = ""; + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new Address(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Address_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Address_fieldAccessorTable + .ensureFieldAccessorsInitialized( + kafka.automq.table.process.proto.PersonProto.Address.class, kafka.automq.table.process.proto.PersonProto.Address.Builder.class); + } + + public static final int STREET_FIELD_NUMBER = 1; + @SuppressWarnings("serial") + private volatile java.lang.Object street_ = ""; + /** + * string street = 1; + * @return The street. + */ + @java.lang.Override + public java.lang.String getStreet() { + java.lang.Object ref = street_; + if (ref instanceof java.lang.String) { + return (java.lang.String) ref; + } else { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + street_ = s; + return s; + } + } + /** + * string street = 1; + * @return The bytes for street. + */ + @java.lang.Override + public com.google.protobuf.ByteString + getStreetBytes() { + java.lang.Object ref = street_; + if (ref instanceof java.lang.String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + street_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + + public static final int CITY_FIELD_NUMBER = 2; + @SuppressWarnings("serial") + private volatile java.lang.Object city_ = ""; + /** + * string city = 2; + * @return The city. + */ + @java.lang.Override + public java.lang.String getCity() { + java.lang.Object ref = city_; + if (ref instanceof java.lang.String) { + return (java.lang.String) ref; + } else { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + city_ = s; + return s; + } + } + /** + * string city = 2; + * @return The bytes for city. + */ + @java.lang.Override + public com.google.protobuf.ByteString + getCityBytes() { + java.lang.Object ref = city_; + if (ref instanceof java.lang.String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + city_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(street_)) { + com.google.protobuf.GeneratedMessageV3.writeString(output, 1, street_); + } + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(city_)) { + com.google.protobuf.GeneratedMessageV3.writeString(output, 2, city_); + } + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(street_)) { + size += com.google.protobuf.GeneratedMessageV3.computeStringSize(1, street_); + } + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(city_)) { + size += com.google.protobuf.GeneratedMessageV3.computeStringSize(2, city_); + } + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof kafka.automq.table.process.proto.PersonProto.Address)) { + return super.equals(obj); + } + kafka.automq.table.process.proto.PersonProto.Address other = (kafka.automq.table.process.proto.PersonProto.Address) obj; + + if (!getStreet() + .equals(other.getStreet())) return false; + if (!getCity() + .equals(other.getCity())) return false; + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (37 * hash) + STREET_FIELD_NUMBER; + hash = (53 * hash) + getStreet().hashCode(); + hash = (37 * hash) + CITY_FIELD_NUMBER; + hash = (53 * hash) + getCity().hashCode(); + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static kafka.automq.table.process.proto.PersonProto.Address parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(kafka.automq.table.process.proto.PersonProto.Address prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code kafka.automq.table.process.proto.Address} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:kafka.automq.table.process.proto.Address) + kafka.automq.table.process.proto.PersonProto.AddressOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Address_descriptor; + } + + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Address_fieldAccessorTable + .ensureFieldAccessorsInitialized( + kafka.automq.table.process.proto.PersonProto.Address.class, kafka.automq.table.process.proto.PersonProto.Address.Builder.class); + } + + // Construct using kafka.automq.table.process.proto.PersonProto.Address.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + bitField0_ = 0; + street_ = ""; + city_ = ""; + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Address_descriptor; + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Address getDefaultInstanceForType() { + return kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance(); + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Address build() { + kafka.automq.table.process.proto.PersonProto.Address result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Address buildPartial() { + kafka.automq.table.process.proto.PersonProto.Address result = new kafka.automq.table.process.proto.PersonProto.Address(this); + if (bitField0_ != 0) { buildPartial0(result); } + onBuilt(); + return result; + } + + private void buildPartial0(kafka.automq.table.process.proto.PersonProto.Address result) { + int from_bitField0_ = bitField0_; + if (((from_bitField0_ & 0x00000001) != 0)) { + result.street_ = street_; + } + if (((from_bitField0_ & 0x00000002) != 0)) { + result.city_ = city_; + } + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof kafka.automq.table.process.proto.PersonProto.Address) { + return mergeFrom((kafka.automq.table.process.proto.PersonProto.Address)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(kafka.automq.table.process.proto.PersonProto.Address other) { + if (other == kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance()) return this; + if (!other.getStreet().isEmpty()) { + street_ = other.street_; + bitField0_ |= 0x00000001; + onChanged(); + } + if (!other.getCity().isEmpty()) { + city_ = other.city_; + bitField0_ |= 0x00000002; + onChanged(); + } + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + case 10: { + street_ = input.readStringRequireUtf8(); + bitField0_ |= 0x00000001; + break; + } // case 10 + case 18: { + city_ = input.readStringRequireUtf8(); + bitField0_ |= 0x00000002; + break; + } // case 18 + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + private int bitField0_; + + private java.lang.Object street_ = ""; + /** + * string street = 1; + * @return The street. + */ + public java.lang.String getStreet() { + java.lang.Object ref = street_; + if (!(ref instanceof java.lang.String)) { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + street_ = s; + return s; + } else { + return (java.lang.String) ref; + } + } + /** + * string street = 1; + * @return The bytes for street. + */ + public com.google.protobuf.ByteString + getStreetBytes() { + java.lang.Object ref = street_; + if (ref instanceof String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + street_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + /** + * string street = 1; + * @param value The street to set. + * @return This builder for chaining. + */ + public Builder setStreet( + java.lang.String value) { + if (value == null) { throw new NullPointerException(); } + street_ = value; + bitField0_ |= 0x00000001; + onChanged(); + return this; + } + /** + * string street = 1; + * @return This builder for chaining. + */ + public Builder clearStreet() { + street_ = getDefaultInstance().getStreet(); + bitField0_ = (bitField0_ & ~0x00000001); + onChanged(); + return this; + } + /** + * string street = 1; + * @param value The bytes for street to set. + * @return This builder for chaining. + */ + public Builder setStreetBytes( + com.google.protobuf.ByteString value) { + if (value == null) { throw new NullPointerException(); } + checkByteStringIsUtf8(value); + street_ = value; + bitField0_ |= 0x00000001; + onChanged(); + return this; + } + + private java.lang.Object city_ = ""; + /** + * string city = 2; + * @return The city. + */ + public java.lang.String getCity() { + java.lang.Object ref = city_; + if (!(ref instanceof java.lang.String)) { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + city_ = s; + return s; + } else { + return (java.lang.String) ref; + } + } + /** + * string city = 2; + * @return The bytes for city. + */ + public com.google.protobuf.ByteString + getCityBytes() { + java.lang.Object ref = city_; + if (ref instanceof String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + city_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + /** + * string city = 2; + * @param value The city to set. + * @return This builder for chaining. + */ + public Builder setCity( + java.lang.String value) { + if (value == null) { throw new NullPointerException(); } + city_ = value; + bitField0_ |= 0x00000002; + onChanged(); + return this; + } + /** + * string city = 2; + * @return This builder for chaining. + */ + public Builder clearCity() { + city_ = getDefaultInstance().getCity(); + bitField0_ = (bitField0_ & ~0x00000002); + onChanged(); + return this; + } + /** + * string city = 2; + * @param value The bytes for city to set. + * @return This builder for chaining. + */ + public Builder setCityBytes( + com.google.protobuf.ByteString value) { + if (value == null) { throw new NullPointerException(); } + checkByteStringIsUtf8(value); + city_ = value; + bitField0_ |= 0x00000002; + onChanged(); + return this; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:kafka.automq.table.process.proto.Address) + } + + // @@protoc_insertion_point(class_scope:kafka.automq.table.process.proto.Address) + private static final kafka.automq.table.process.proto.PersonProto.Address DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new kafka.automq.table.process.proto.PersonProto.Address(); + } + + public static kafka.automq.table.process.proto.PersonProto.Address getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser
    + PARSER = new com.google.protobuf.AbstractParser
    () { + @java.lang.Override + public Address parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser
    parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser
    getParserForType() { + return PARSER; + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Address getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + public interface PersonOrBuilder extends + // @@protoc_insertion_point(interface_extends:kafka.automq.table.process.proto.Person) + com.google.protobuf.MessageOrBuilder { + + /** + * int64 id = 1; + * @return The id. + */ + long getId(); + + /** + * string name = 2; + * @return The name. + */ + java.lang.String getName(); + /** + * string name = 2; + * @return The bytes for name. + */ + com.google.protobuf.ByteString + getNameBytes(); + + /** + * bool is_active = 3; + * @return The isActive. + */ + boolean getIsActive(); + + /** + * bytes extra_data = 4; + * @return The extraData. + */ + com.google.protobuf.ByteString getExtraData(); + + /** + * .kafka.automq.table.process.proto.Address address = 5; + * @return Whether the address field is set. + */ + boolean hasAddress(); + /** + * .kafka.automq.table.process.proto.Address address = 5; + * @return The address. + */ + kafka.automq.table.process.proto.PersonProto.Address getAddress(); + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + kafka.automq.table.process.proto.PersonProto.AddressOrBuilder getAddressOrBuilder(); + + /** + * repeated string roles = 6; + * @return A list containing the roles. + */ + java.util.List + getRolesList(); + /** + * repeated string roles = 6; + * @return The count of roles. + */ + int getRolesCount(); + /** + * repeated string roles = 6; + * @param index The index of the element to return. + * @return The roles at the given index. + */ + java.lang.String getRoles(int index); + /** + * repeated string roles = 6; + * @param index The index of the value to return. + * @return The bytes of the roles at the given index. + */ + com.google.protobuf.ByteString + getRolesBytes(int index); + + /** + * map<string, string> attributes = 7; + */ + int getAttributesCount(); + /** + * map<string, string> attributes = 7; + */ + boolean containsAttributes( + java.lang.String key); + /** + * Use {@link #getAttributesMap()} instead. + */ + @java.lang.Deprecated + java.util.Map + getAttributes(); + /** + * map<string, string> attributes = 7; + */ + java.util.Map + getAttributesMap(); + /** + * map<string, string> attributes = 7; + */ + /* nullable */ +java.lang.String getAttributesOrDefault( + java.lang.String key, + /* nullable */ +java.lang.String defaultValue); + /** + * map<string, string> attributes = 7; + */ + java.lang.String getAttributesOrThrow( + java.lang.String key); + + /** + * .google.protobuf.Timestamp last_updated = 8; + * @return Whether the lastUpdated field is set. + */ + boolean hasLastUpdated(); + /** + * .google.protobuf.Timestamp last_updated = 8; + * @return The lastUpdated. + */ + com.google.protobuf.Timestamp getLastUpdated(); + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + com.google.protobuf.TimestampOrBuilder getLastUpdatedOrBuilder(); + } + /** + * Protobuf type {@code kafka.automq.table.process.proto.Person} + */ + public static final class Person extends + com.google.protobuf.GeneratedMessageV3 implements + // @@protoc_insertion_point(message_implements:kafka.automq.table.process.proto.Person) + PersonOrBuilder { + private static final long serialVersionUID = 0L; + // Use Person.newBuilder() to construct. + private Person(com.google.protobuf.GeneratedMessageV3.Builder builder) { + super(builder); + } + private Person() { + name_ = ""; + extraData_ = com.google.protobuf.ByteString.EMPTY; + roles_ = com.google.protobuf.LazyStringArrayList.EMPTY; + } + + @java.lang.Override + @SuppressWarnings({"unused"}) + protected java.lang.Object newInstance( + UnusedPrivateParameter unused) { + return new Person(); + } + + @java.lang.Override + public final com.google.protobuf.UnknownFieldSet + getUnknownFields() { + return this.unknownFields; + } + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Person_descriptor; + } + + @SuppressWarnings({"rawtypes"}) + @java.lang.Override + protected com.google.protobuf.MapField internalGetMapField( + int number) { + switch (number) { + case 7: + return internalGetAttributes(); + default: + throw new RuntimeException( + "Invalid map field number: " + number); + } + } + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Person_fieldAccessorTable + .ensureFieldAccessorsInitialized( + kafka.automq.table.process.proto.PersonProto.Person.class, kafka.automq.table.process.proto.PersonProto.Person.Builder.class); + } + + public static final int ID_FIELD_NUMBER = 1; + private long id_ = 0L; + /** + * int64 id = 1; + * @return The id. + */ + @java.lang.Override + public long getId() { + return id_; + } + + public static final int NAME_FIELD_NUMBER = 2; + @SuppressWarnings("serial") + private volatile java.lang.Object name_ = ""; + /** + * string name = 2; + * @return The name. + */ + @java.lang.Override + public java.lang.String getName() { + java.lang.Object ref = name_; + if (ref instanceof java.lang.String) { + return (java.lang.String) ref; + } else { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + name_ = s; + return s; + } + } + /** + * string name = 2; + * @return The bytes for name. + */ + @java.lang.Override + public com.google.protobuf.ByteString + getNameBytes() { + java.lang.Object ref = name_; + if (ref instanceof java.lang.String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + name_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + + public static final int IS_ACTIVE_FIELD_NUMBER = 3; + private boolean isActive_ = false; + /** + * bool is_active = 3; + * @return The isActive. + */ + @java.lang.Override + public boolean getIsActive() { + return isActive_; + } + + public static final int EXTRA_DATA_FIELD_NUMBER = 4; + private com.google.protobuf.ByteString extraData_ = com.google.protobuf.ByteString.EMPTY; + /** + * bytes extra_data = 4; + * @return The extraData. + */ + @java.lang.Override + public com.google.protobuf.ByteString getExtraData() { + return extraData_; + } + + public static final int ADDRESS_FIELD_NUMBER = 5; + private kafka.automq.table.process.proto.PersonProto.Address address_; + /** + * .kafka.automq.table.process.proto.Address address = 5; + * @return Whether the address field is set. + */ + @java.lang.Override + public boolean hasAddress() { + return address_ != null; + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + * @return The address. + */ + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Address getAddress() { + return address_ == null ? kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance() : address_; + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.AddressOrBuilder getAddressOrBuilder() { + return address_ == null ? kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance() : address_; + } + + public static final int ROLES_FIELD_NUMBER = 6; + @SuppressWarnings("serial") + private com.google.protobuf.LazyStringList roles_; + /** + * repeated string roles = 6; + * @return A list containing the roles. + */ + public com.google.protobuf.ProtocolStringList + getRolesList() { + return roles_; + } + /** + * repeated string roles = 6; + * @return The count of roles. + */ + public int getRolesCount() { + return roles_.size(); + } + /** + * repeated string roles = 6; + * @param index The index of the element to return. + * @return The roles at the given index. + */ + public java.lang.String getRoles(int index) { + return roles_.get(index); + } + /** + * repeated string roles = 6; + * @param index The index of the value to return. + * @return The bytes of the roles at the given index. + */ + public com.google.protobuf.ByteString + getRolesBytes(int index) { + return roles_.getByteString(index); + } + + public static final int ATTRIBUTES_FIELD_NUMBER = 7; + private static final class AttributesDefaultEntryHolder { + static final com.google.protobuf.MapEntry< + java.lang.String, java.lang.String> defaultEntry = + com.google.protobuf.MapEntry + .newDefaultInstance( + kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Person_AttributesEntry_descriptor, + com.google.protobuf.WireFormat.FieldType.STRING, + "", + com.google.protobuf.WireFormat.FieldType.STRING, + ""); + } + @SuppressWarnings("serial") + private com.google.protobuf.MapField< + java.lang.String, java.lang.String> attributes_; + private com.google.protobuf.MapField + internalGetAttributes() { + if (attributes_ == null) { + return com.google.protobuf.MapField.emptyMapField( + AttributesDefaultEntryHolder.defaultEntry); + } + return attributes_; + } + public int getAttributesCount() { + return internalGetAttributes().getMap().size(); + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public boolean containsAttributes( + java.lang.String key) { + if (key == null) { throw new NullPointerException("map key"); } + return internalGetAttributes().getMap().containsKey(key); + } + /** + * Use {@link #getAttributesMap()} instead. + */ + @java.lang.Override + @java.lang.Deprecated + public java.util.Map getAttributes() { + return getAttributesMap(); + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public java.util.Map getAttributesMap() { + return internalGetAttributes().getMap(); + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public /* nullable */ +java.lang.String getAttributesOrDefault( + java.lang.String key, + /* nullable */ +java.lang.String defaultValue) { + if (key == null) { throw new NullPointerException("map key"); } + java.util.Map map = + internalGetAttributes().getMap(); + return map.containsKey(key) ? map.get(key) : defaultValue; + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public java.lang.String getAttributesOrThrow( + java.lang.String key) { + if (key == null) { throw new NullPointerException("map key"); } + java.util.Map map = + internalGetAttributes().getMap(); + if (!map.containsKey(key)) { + throw new java.lang.IllegalArgumentException(); + } + return map.get(key); + } + + public static final int LAST_UPDATED_FIELD_NUMBER = 8; + private com.google.protobuf.Timestamp lastUpdated_; + /** + * .google.protobuf.Timestamp last_updated = 8; + * @return Whether the lastUpdated field is set. + */ + @java.lang.Override + public boolean hasLastUpdated() { + return lastUpdated_ != null; + } + /** + * .google.protobuf.Timestamp last_updated = 8; + * @return The lastUpdated. + */ + @java.lang.Override + public com.google.protobuf.Timestamp getLastUpdated() { + return lastUpdated_ == null ? com.google.protobuf.Timestamp.getDefaultInstance() : lastUpdated_; + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + @java.lang.Override + public com.google.protobuf.TimestampOrBuilder getLastUpdatedOrBuilder() { + return lastUpdated_ == null ? com.google.protobuf.Timestamp.getDefaultInstance() : lastUpdated_; + } + + private byte memoizedIsInitialized = -1; + @java.lang.Override + public final boolean isInitialized() { + byte isInitialized = memoizedIsInitialized; + if (isInitialized == 1) return true; + if (isInitialized == 0) return false; + + memoizedIsInitialized = 1; + return true; + } + + @java.lang.Override + public void writeTo(com.google.protobuf.CodedOutputStream output) + throws java.io.IOException { + if (id_ != 0L) { + output.writeInt64(1, id_); + } + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(name_)) { + com.google.protobuf.GeneratedMessageV3.writeString(output, 2, name_); + } + if (isActive_ != false) { + output.writeBool(3, isActive_); + } + if (!extraData_.isEmpty()) { + output.writeBytes(4, extraData_); + } + if (address_ != null) { + output.writeMessage(5, getAddress()); + } + for (int i = 0; i < roles_.size(); i++) { + com.google.protobuf.GeneratedMessageV3.writeString(output, 6, roles_.getRaw(i)); + } + com.google.protobuf.GeneratedMessageV3 + .serializeStringMapTo( + output, + internalGetAttributes(), + AttributesDefaultEntryHolder.defaultEntry, + 7); + if (lastUpdated_ != null) { + output.writeMessage(8, getLastUpdated()); + } + getUnknownFields().writeTo(output); + } + + @java.lang.Override + public int getSerializedSize() { + int size = memoizedSize; + if (size != -1) return size; + + size = 0; + if (id_ != 0L) { + size += com.google.protobuf.CodedOutputStream + .computeInt64Size(1, id_); + } + if (!com.google.protobuf.GeneratedMessageV3.isStringEmpty(name_)) { + size += com.google.protobuf.GeneratedMessageV3.computeStringSize(2, name_); + } + if (isActive_ != false) { + size += com.google.protobuf.CodedOutputStream + .computeBoolSize(3, isActive_); + } + if (!extraData_.isEmpty()) { + size += com.google.protobuf.CodedOutputStream + .computeBytesSize(4, extraData_); + } + if (address_ != null) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(5, getAddress()); + } + { + int dataSize = 0; + for (int i = 0; i < roles_.size(); i++) { + dataSize += computeStringSizeNoTag(roles_.getRaw(i)); + } + size += dataSize; + size += 1 * getRolesList().size(); + } + for (java.util.Map.Entry entry + : internalGetAttributes().getMap().entrySet()) { + com.google.protobuf.MapEntry + attributes__ = AttributesDefaultEntryHolder.defaultEntry.newBuilderForType() + .setKey(entry.getKey()) + .setValue(entry.getValue()) + .build(); + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(7, attributes__); + } + if (lastUpdated_ != null) { + size += com.google.protobuf.CodedOutputStream + .computeMessageSize(8, getLastUpdated()); + } + size += getUnknownFields().getSerializedSize(); + memoizedSize = size; + return size; + } + + @java.lang.Override + public boolean equals(final java.lang.Object obj) { + if (obj == this) { + return true; + } + if (!(obj instanceof kafka.automq.table.process.proto.PersonProto.Person)) { + return super.equals(obj); + } + kafka.automq.table.process.proto.PersonProto.Person other = (kafka.automq.table.process.proto.PersonProto.Person) obj; + + if (getId() + != other.getId()) return false; + if (!getName() + .equals(other.getName())) return false; + if (getIsActive() + != other.getIsActive()) return false; + if (!getExtraData() + .equals(other.getExtraData())) return false; + if (hasAddress() != other.hasAddress()) return false; + if (hasAddress()) { + if (!getAddress() + .equals(other.getAddress())) return false; + } + if (!getRolesList() + .equals(other.getRolesList())) return false; + if (!internalGetAttributes().equals( + other.internalGetAttributes())) return false; + if (hasLastUpdated() != other.hasLastUpdated()) return false; + if (hasLastUpdated()) { + if (!getLastUpdated() + .equals(other.getLastUpdated())) return false; + } + if (!getUnknownFields().equals(other.getUnknownFields())) return false; + return true; + } + + @java.lang.Override + public int hashCode() { + if (memoizedHashCode != 0) { + return memoizedHashCode; + } + int hash = 41; + hash = (19 * hash) + getDescriptor().hashCode(); + hash = (37 * hash) + ID_FIELD_NUMBER; + hash = (53 * hash) + com.google.protobuf.Internal.hashLong( + getId()); + hash = (37 * hash) + NAME_FIELD_NUMBER; + hash = (53 * hash) + getName().hashCode(); + hash = (37 * hash) + IS_ACTIVE_FIELD_NUMBER; + hash = (53 * hash) + com.google.protobuf.Internal.hashBoolean( + getIsActive()); + hash = (37 * hash) + EXTRA_DATA_FIELD_NUMBER; + hash = (53 * hash) + getExtraData().hashCode(); + if (hasAddress()) { + hash = (37 * hash) + ADDRESS_FIELD_NUMBER; + hash = (53 * hash) + getAddress().hashCode(); + } + if (getRolesCount() > 0) { + hash = (37 * hash) + ROLES_FIELD_NUMBER; + hash = (53 * hash) + getRolesList().hashCode(); + } + if (!internalGetAttributes().getMap().isEmpty()) { + hash = (37 * hash) + ATTRIBUTES_FIELD_NUMBER; + hash = (53 * hash) + internalGetAttributes().hashCode(); + } + if (hasLastUpdated()) { + hash = (37 * hash) + LAST_UPDATED_FIELD_NUMBER; + hash = (53 * hash) + getLastUpdated().hashCode(); + } + hash = (29 * hash) + getUnknownFields().hashCode(); + memoizedHashCode = hash; + return hash; + } + + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + java.nio.ByteBuffer data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + java.nio.ByteBuffer data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + com.google.protobuf.ByteString data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + com.google.protobuf.ByteString data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom(byte[] data) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + byte[] data, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + return PARSER.parseFrom(data, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseDelimitedFrom(java.io.InputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseDelimitedFrom( + java.io.InputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseDelimitedWithIOException(PARSER, input, extensionRegistry); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + com.google.protobuf.CodedInputStream input) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input); + } + public static kafka.automq.table.process.proto.PersonProto.Person parseFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + return com.google.protobuf.GeneratedMessageV3 + .parseWithIOException(PARSER, input, extensionRegistry); + } + + @java.lang.Override + public Builder newBuilderForType() { return newBuilder(); } + public static Builder newBuilder() { + return DEFAULT_INSTANCE.toBuilder(); + } + public static Builder newBuilder(kafka.automq.table.process.proto.PersonProto.Person prototype) { + return DEFAULT_INSTANCE.toBuilder().mergeFrom(prototype); + } + @java.lang.Override + public Builder toBuilder() { + return this == DEFAULT_INSTANCE + ? new Builder() : new Builder().mergeFrom(this); + } + + @java.lang.Override + protected Builder newBuilderForType( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + Builder builder = new Builder(parent); + return builder; + } + /** + * Protobuf type {@code kafka.automq.table.process.proto.Person} + */ + public static final class Builder extends + com.google.protobuf.GeneratedMessageV3.Builder implements + // @@protoc_insertion_point(builder_implements:kafka.automq.table.process.proto.Person) + kafka.automq.table.process.proto.PersonProto.PersonOrBuilder { + public static final com.google.protobuf.Descriptors.Descriptor + getDescriptor() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Person_descriptor; + } + + @SuppressWarnings({"rawtypes"}) + protected com.google.protobuf.MapField internalGetMapField( + int number) { + switch (number) { + case 7: + return internalGetAttributes(); + default: + throw new RuntimeException( + "Invalid map field number: " + number); + } + } + @SuppressWarnings({"rawtypes"}) + protected com.google.protobuf.MapField internalGetMutableMapField( + int number) { + switch (number) { + case 7: + return internalGetMutableAttributes(); + default: + throw new RuntimeException( + "Invalid map field number: " + number); + } + } + @java.lang.Override + protected com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internalGetFieldAccessorTable() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Person_fieldAccessorTable + .ensureFieldAccessorsInitialized( + kafka.automq.table.process.proto.PersonProto.Person.class, kafka.automq.table.process.proto.PersonProto.Person.Builder.class); + } + + // Construct using kafka.automq.table.process.proto.PersonProto.Person.newBuilder() + private Builder() { + + } + + private Builder( + com.google.protobuf.GeneratedMessageV3.BuilderParent parent) { + super(parent); + + } + @java.lang.Override + public Builder clear() { + super.clear(); + bitField0_ = 0; + id_ = 0L; + name_ = ""; + isActive_ = false; + extraData_ = com.google.protobuf.ByteString.EMPTY; + address_ = null; + if (addressBuilder_ != null) { + addressBuilder_.dispose(); + addressBuilder_ = null; + } + roles_ = com.google.protobuf.LazyStringArrayList.EMPTY; + bitField0_ = (bitField0_ & ~0x00000020); + internalGetMutableAttributes().clear(); + lastUpdated_ = null; + if (lastUpdatedBuilder_ != null) { + lastUpdatedBuilder_.dispose(); + lastUpdatedBuilder_ = null; + } + return this; + } + + @java.lang.Override + public com.google.protobuf.Descriptors.Descriptor + getDescriptorForType() { + return kafka.automq.table.process.proto.PersonProto.internal_static_kafka_automq_table_process_proto_Person_descriptor; + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Person getDefaultInstanceForType() { + return kafka.automq.table.process.proto.PersonProto.Person.getDefaultInstance(); + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Person build() { + kafka.automq.table.process.proto.PersonProto.Person result = buildPartial(); + if (!result.isInitialized()) { + throw newUninitializedMessageException(result); + } + return result; + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Person buildPartial() { + kafka.automq.table.process.proto.PersonProto.Person result = new kafka.automq.table.process.proto.PersonProto.Person(this); + buildPartialRepeatedFields(result); + if (bitField0_ != 0) { buildPartial0(result); } + onBuilt(); + return result; + } + + private void buildPartialRepeatedFields(kafka.automq.table.process.proto.PersonProto.Person result) { + if (((bitField0_ & 0x00000020) != 0)) { + roles_ = roles_.getUnmodifiableView(); + bitField0_ = (bitField0_ & ~0x00000020); + } + result.roles_ = roles_; + } + + private void buildPartial0(kafka.automq.table.process.proto.PersonProto.Person result) { + int from_bitField0_ = bitField0_; + if (((from_bitField0_ & 0x00000001) != 0)) { + result.id_ = id_; + } + if (((from_bitField0_ & 0x00000002) != 0)) { + result.name_ = name_; + } + if (((from_bitField0_ & 0x00000004) != 0)) { + result.isActive_ = isActive_; + } + if (((from_bitField0_ & 0x00000008) != 0)) { + result.extraData_ = extraData_; + } + if (((from_bitField0_ & 0x00000010) != 0)) { + result.address_ = addressBuilder_ == null + ? address_ + : addressBuilder_.build(); + } + if (((from_bitField0_ & 0x00000040) != 0)) { + result.attributes_ = internalGetAttributes(); + result.attributes_.makeImmutable(); + } + if (((from_bitField0_ & 0x00000080) != 0)) { + result.lastUpdated_ = lastUpdatedBuilder_ == null + ? lastUpdated_ + : lastUpdatedBuilder_.build(); + } + } + + @java.lang.Override + public Builder clone() { + return super.clone(); + } + @java.lang.Override + public Builder setField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.setField(field, value); + } + @java.lang.Override + public Builder clearField( + com.google.protobuf.Descriptors.FieldDescriptor field) { + return super.clearField(field); + } + @java.lang.Override + public Builder clearOneof( + com.google.protobuf.Descriptors.OneofDescriptor oneof) { + return super.clearOneof(oneof); + } + @java.lang.Override + public Builder setRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + int index, java.lang.Object value) { + return super.setRepeatedField(field, index, value); + } + @java.lang.Override + public Builder addRepeatedField( + com.google.protobuf.Descriptors.FieldDescriptor field, + java.lang.Object value) { + return super.addRepeatedField(field, value); + } + @java.lang.Override + public Builder mergeFrom(com.google.protobuf.Message other) { + if (other instanceof kafka.automq.table.process.proto.PersonProto.Person) { + return mergeFrom((kafka.automq.table.process.proto.PersonProto.Person)other); + } else { + super.mergeFrom(other); + return this; + } + } + + public Builder mergeFrom(kafka.automq.table.process.proto.PersonProto.Person other) { + if (other == kafka.automq.table.process.proto.PersonProto.Person.getDefaultInstance()) return this; + if (other.getId() != 0L) { + setId(other.getId()); + } + if (!other.getName().isEmpty()) { + name_ = other.name_; + bitField0_ |= 0x00000002; + onChanged(); + } + if (other.getIsActive() != false) { + setIsActive(other.getIsActive()); + } + if (other.getExtraData() != com.google.protobuf.ByteString.EMPTY) { + setExtraData(other.getExtraData()); + } + if (other.hasAddress()) { + mergeAddress(other.getAddress()); + } + if (!other.roles_.isEmpty()) { + if (roles_.isEmpty()) { + roles_ = other.roles_; + bitField0_ = (bitField0_ & ~0x00000020); + } else { + ensureRolesIsMutable(); + roles_.addAll(other.roles_); + } + onChanged(); + } + internalGetMutableAttributes().mergeFrom( + other.internalGetAttributes()); + bitField0_ |= 0x00000040; + if (other.hasLastUpdated()) { + mergeLastUpdated(other.getLastUpdated()); + } + this.mergeUnknownFields(other.getUnknownFields()); + onChanged(); + return this; + } + + @java.lang.Override + public final boolean isInitialized() { + return true; + } + + @java.lang.Override + public Builder mergeFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws java.io.IOException { + if (extensionRegistry == null) { + throw new java.lang.NullPointerException(); + } + try { + boolean done = false; + while (!done) { + int tag = input.readTag(); + switch (tag) { + case 0: + done = true; + break; + case 8: { + id_ = input.readInt64(); + bitField0_ |= 0x00000001; + break; + } // case 8 + case 18: { + name_ = input.readStringRequireUtf8(); + bitField0_ |= 0x00000002; + break; + } // case 18 + case 24: { + isActive_ = input.readBool(); + bitField0_ |= 0x00000004; + break; + } // case 24 + case 34: { + extraData_ = input.readBytes(); + bitField0_ |= 0x00000008; + break; + } // case 34 + case 42: { + input.readMessage( + getAddressFieldBuilder().getBuilder(), + extensionRegistry); + bitField0_ |= 0x00000010; + break; + } // case 42 + case 50: { + java.lang.String s = input.readStringRequireUtf8(); + ensureRolesIsMutable(); + roles_.add(s); + break; + } // case 50 + case 58: { + com.google.protobuf.MapEntry + attributes__ = input.readMessage( + AttributesDefaultEntryHolder.defaultEntry.getParserForType(), extensionRegistry); + internalGetMutableAttributes().getMutableMap().put( + attributes__.getKey(), attributes__.getValue()); + bitField0_ |= 0x00000040; + break; + } // case 58 + case 66: { + input.readMessage( + getLastUpdatedFieldBuilder().getBuilder(), + extensionRegistry); + bitField0_ |= 0x00000080; + break; + } // case 66 + default: { + if (!super.parseUnknownField(input, extensionRegistry, tag)) { + done = true; // was an endgroup tag + } + break; + } // default: + } // switch (tag) + } // while (!done) + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.unwrapIOException(); + } finally { + onChanged(); + } // finally + return this; + } + private int bitField0_; + + private long id_ ; + /** + * int64 id = 1; + * @return The id. + */ + @java.lang.Override + public long getId() { + return id_; + } + /** + * int64 id = 1; + * @param value The id to set. + * @return This builder for chaining. + */ + public Builder setId(long value) { + + id_ = value; + bitField0_ |= 0x00000001; + onChanged(); + return this; + } + /** + * int64 id = 1; + * @return This builder for chaining. + */ + public Builder clearId() { + bitField0_ = (bitField0_ & ~0x00000001); + id_ = 0L; + onChanged(); + return this; + } + + private java.lang.Object name_ = ""; + /** + * string name = 2; + * @return The name. + */ + public java.lang.String getName() { + java.lang.Object ref = name_; + if (!(ref instanceof java.lang.String)) { + com.google.protobuf.ByteString bs = + (com.google.protobuf.ByteString) ref; + java.lang.String s = bs.toStringUtf8(); + name_ = s; + return s; + } else { + return (java.lang.String) ref; + } + } + /** + * string name = 2; + * @return The bytes for name. + */ + public com.google.protobuf.ByteString + getNameBytes() { + java.lang.Object ref = name_; + if (ref instanceof String) { + com.google.protobuf.ByteString b = + com.google.protobuf.ByteString.copyFromUtf8( + (java.lang.String) ref); + name_ = b; + return b; + } else { + return (com.google.protobuf.ByteString) ref; + } + } + /** + * string name = 2; + * @param value The name to set. + * @return This builder for chaining. + */ + public Builder setName( + java.lang.String value) { + if (value == null) { throw new NullPointerException(); } + name_ = value; + bitField0_ |= 0x00000002; + onChanged(); + return this; + } + /** + * string name = 2; + * @return This builder for chaining. + */ + public Builder clearName() { + name_ = getDefaultInstance().getName(); + bitField0_ = (bitField0_ & ~0x00000002); + onChanged(); + return this; + } + /** + * string name = 2; + * @param value The bytes for name to set. + * @return This builder for chaining. + */ + public Builder setNameBytes( + com.google.protobuf.ByteString value) { + if (value == null) { throw new NullPointerException(); } + checkByteStringIsUtf8(value); + name_ = value; + bitField0_ |= 0x00000002; + onChanged(); + return this; + } + + private boolean isActive_ ; + /** + * bool is_active = 3; + * @return The isActive. + */ + @java.lang.Override + public boolean getIsActive() { + return isActive_; + } + /** + * bool is_active = 3; + * @param value The isActive to set. + * @return This builder for chaining. + */ + public Builder setIsActive(boolean value) { + + isActive_ = value; + bitField0_ |= 0x00000004; + onChanged(); + return this; + } + /** + * bool is_active = 3; + * @return This builder for chaining. + */ + public Builder clearIsActive() { + bitField0_ = (bitField0_ & ~0x00000004); + isActive_ = false; + onChanged(); + return this; + } + + private com.google.protobuf.ByteString extraData_ = com.google.protobuf.ByteString.EMPTY; + /** + * bytes extra_data = 4; + * @return The extraData. + */ + @java.lang.Override + public com.google.protobuf.ByteString getExtraData() { + return extraData_; + } + /** + * bytes extra_data = 4; + * @param value The extraData to set. + * @return This builder for chaining. + */ + public Builder setExtraData(com.google.protobuf.ByteString value) { + if (value == null) { throw new NullPointerException(); } + extraData_ = value; + bitField0_ |= 0x00000008; + onChanged(); + return this; + } + /** + * bytes extra_data = 4; + * @return This builder for chaining. + */ + public Builder clearExtraData() { + bitField0_ = (bitField0_ & ~0x00000008); + extraData_ = getDefaultInstance().getExtraData(); + onChanged(); + return this; + } + + private kafka.automq.table.process.proto.PersonProto.Address address_; + private com.google.protobuf.SingleFieldBuilderV3< + kafka.automq.table.process.proto.PersonProto.Address, kafka.automq.table.process.proto.PersonProto.Address.Builder, kafka.automq.table.process.proto.PersonProto.AddressOrBuilder> addressBuilder_; + /** + * .kafka.automq.table.process.proto.Address address = 5; + * @return Whether the address field is set. + */ + public boolean hasAddress() { + return ((bitField0_ & 0x00000010) != 0); + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + * @return The address. + */ + public kafka.automq.table.process.proto.PersonProto.Address getAddress() { + if (addressBuilder_ == null) { + return address_ == null ? kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance() : address_; + } else { + return addressBuilder_.getMessage(); + } + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + public Builder setAddress(kafka.automq.table.process.proto.PersonProto.Address value) { + if (addressBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + address_ = value; + } else { + addressBuilder_.setMessage(value); + } + bitField0_ |= 0x00000010; + onChanged(); + return this; + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + public Builder setAddress( + kafka.automq.table.process.proto.PersonProto.Address.Builder builderForValue) { + if (addressBuilder_ == null) { + address_ = builderForValue.build(); + } else { + addressBuilder_.setMessage(builderForValue.build()); + } + bitField0_ |= 0x00000010; + onChanged(); + return this; + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + public Builder mergeAddress(kafka.automq.table.process.proto.PersonProto.Address value) { + if (addressBuilder_ == null) { + if (((bitField0_ & 0x00000010) != 0) && + address_ != null && + address_ != kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance()) { + getAddressBuilder().mergeFrom(value); + } else { + address_ = value; + } + } else { + addressBuilder_.mergeFrom(value); + } + bitField0_ |= 0x00000010; + onChanged(); + return this; + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + public Builder clearAddress() { + bitField0_ = (bitField0_ & ~0x00000010); + address_ = null; + if (addressBuilder_ != null) { + addressBuilder_.dispose(); + addressBuilder_ = null; + } + onChanged(); + return this; + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + public kafka.automq.table.process.proto.PersonProto.Address.Builder getAddressBuilder() { + bitField0_ |= 0x00000010; + onChanged(); + return getAddressFieldBuilder().getBuilder(); + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + public kafka.automq.table.process.proto.PersonProto.AddressOrBuilder getAddressOrBuilder() { + if (addressBuilder_ != null) { + return addressBuilder_.getMessageOrBuilder(); + } else { + return address_ == null ? + kafka.automq.table.process.proto.PersonProto.Address.getDefaultInstance() : address_; + } + } + /** + * .kafka.automq.table.process.proto.Address address = 5; + */ + private com.google.protobuf.SingleFieldBuilderV3< + kafka.automq.table.process.proto.PersonProto.Address, kafka.automq.table.process.proto.PersonProto.Address.Builder, kafka.automq.table.process.proto.PersonProto.AddressOrBuilder> + getAddressFieldBuilder() { + if (addressBuilder_ == null) { + addressBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + kafka.automq.table.process.proto.PersonProto.Address, kafka.automq.table.process.proto.PersonProto.Address.Builder, kafka.automq.table.process.proto.PersonProto.AddressOrBuilder>( + getAddress(), + getParentForChildren(), + isClean()); + address_ = null; + } + return addressBuilder_; + } + + private com.google.protobuf.LazyStringList roles_ = com.google.protobuf.LazyStringArrayList.EMPTY; + private void ensureRolesIsMutable() { + if (!((bitField0_ & 0x00000020) != 0)) { + roles_ = new com.google.protobuf.LazyStringArrayList(roles_); + bitField0_ |= 0x00000020; + } + } + /** + * repeated string roles = 6; + * @return A list containing the roles. + */ + public com.google.protobuf.ProtocolStringList + getRolesList() { + return roles_.getUnmodifiableView(); + } + /** + * repeated string roles = 6; + * @return The count of roles. + */ + public int getRolesCount() { + return roles_.size(); + } + /** + * repeated string roles = 6; + * @param index The index of the element to return. + * @return The roles at the given index. + */ + public java.lang.String getRoles(int index) { + return roles_.get(index); + } + /** + * repeated string roles = 6; + * @param index The index of the value to return. + * @return The bytes of the roles at the given index. + */ + public com.google.protobuf.ByteString + getRolesBytes(int index) { + return roles_.getByteString(index); + } + /** + * repeated string roles = 6; + * @param index The index to set the value at. + * @param value The roles to set. + * @return This builder for chaining. + */ + public Builder setRoles( + int index, java.lang.String value) { + if (value == null) { throw new NullPointerException(); } + ensureRolesIsMutable(); + roles_.set(index, value); + onChanged(); + return this; + } + /** + * repeated string roles = 6; + * @param value The roles to add. + * @return This builder for chaining. + */ + public Builder addRoles( + java.lang.String value) { + if (value == null) { throw new NullPointerException(); } + ensureRolesIsMutable(); + roles_.add(value); + onChanged(); + return this; + } + /** + * repeated string roles = 6; + * @param values The roles to add. + * @return This builder for chaining. + */ + public Builder addAllRoles( + java.lang.Iterable values) { + ensureRolesIsMutable(); + com.google.protobuf.AbstractMessageLite.Builder.addAll( + values, roles_); + onChanged(); + return this; + } + /** + * repeated string roles = 6; + * @return This builder for chaining. + */ + public Builder clearRoles() { + roles_ = com.google.protobuf.LazyStringArrayList.EMPTY; + bitField0_ = (bitField0_ & ~0x00000020); + onChanged(); + return this; + } + /** + * repeated string roles = 6; + * @param value The bytes of the roles to add. + * @return This builder for chaining. + */ + public Builder addRolesBytes( + com.google.protobuf.ByteString value) { + if (value == null) { throw new NullPointerException(); } + checkByteStringIsUtf8(value); + ensureRolesIsMutable(); + roles_.add(value); + onChanged(); + return this; + } + + private com.google.protobuf.MapField< + java.lang.String, java.lang.String> attributes_; + private com.google.protobuf.MapField + internalGetAttributes() { + if (attributes_ == null) { + return com.google.protobuf.MapField.emptyMapField( + AttributesDefaultEntryHolder.defaultEntry); + } + return attributes_; + } + private com.google.protobuf.MapField + internalGetMutableAttributes() { + if (attributes_ == null) { + attributes_ = com.google.protobuf.MapField.newMapField( + AttributesDefaultEntryHolder.defaultEntry); + } + if (!attributes_.isMutable()) { + attributes_ = attributes_.copy(); + } + bitField0_ |= 0x00000040; + onChanged(); + return attributes_; + } + public int getAttributesCount() { + return internalGetAttributes().getMap().size(); + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public boolean containsAttributes( + java.lang.String key) { + if (key == null) { throw new NullPointerException("map key"); } + return internalGetAttributes().getMap().containsKey(key); + } + /** + * Use {@link #getAttributesMap()} instead. + */ + @java.lang.Override + @java.lang.Deprecated + public java.util.Map getAttributes() { + return getAttributesMap(); + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public java.util.Map getAttributesMap() { + return internalGetAttributes().getMap(); + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public /* nullable */ +java.lang.String getAttributesOrDefault( + java.lang.String key, + /* nullable */ +java.lang.String defaultValue) { + if (key == null) { throw new NullPointerException("map key"); } + java.util.Map map = + internalGetAttributes().getMap(); + return map.containsKey(key) ? map.get(key) : defaultValue; + } + /** + * map<string, string> attributes = 7; + */ + @java.lang.Override + public java.lang.String getAttributesOrThrow( + java.lang.String key) { + if (key == null) { throw new NullPointerException("map key"); } + java.util.Map map = + internalGetAttributes().getMap(); + if (!map.containsKey(key)) { + throw new java.lang.IllegalArgumentException(); + } + return map.get(key); + } + public Builder clearAttributes() { + bitField0_ = (bitField0_ & ~0x00000040); + internalGetMutableAttributes().getMutableMap() + .clear(); + return this; + } + /** + * map<string, string> attributes = 7; + */ + public Builder removeAttributes( + java.lang.String key) { + if (key == null) { throw new NullPointerException("map key"); } + internalGetMutableAttributes().getMutableMap() + .remove(key); + return this; + } + /** + * Use alternate mutation accessors instead. + */ + @java.lang.Deprecated + public java.util.Map + getMutableAttributes() { + bitField0_ |= 0x00000040; + return internalGetMutableAttributes().getMutableMap(); + } + /** + * map<string, string> attributes = 7; + */ + public Builder putAttributes( + java.lang.String key, + java.lang.String value) { + if (key == null) { throw new NullPointerException("map key"); } + if (value == null) { throw new NullPointerException("map value"); } + internalGetMutableAttributes().getMutableMap() + .put(key, value); + bitField0_ |= 0x00000040; + return this; + } + /** + * map<string, string> attributes = 7; + */ + public Builder putAllAttributes( + java.util.Map values) { + internalGetMutableAttributes().getMutableMap() + .putAll(values); + bitField0_ |= 0x00000040; + return this; + } + + private com.google.protobuf.Timestamp lastUpdated_; + private com.google.protobuf.SingleFieldBuilderV3< + com.google.protobuf.Timestamp, com.google.protobuf.Timestamp.Builder, com.google.protobuf.TimestampOrBuilder> lastUpdatedBuilder_; + /** + * .google.protobuf.Timestamp last_updated = 8; + * @return Whether the lastUpdated field is set. + */ + public boolean hasLastUpdated() { + return ((bitField0_ & 0x00000080) != 0); + } + /** + * .google.protobuf.Timestamp last_updated = 8; + * @return The lastUpdated. + */ + public com.google.protobuf.Timestamp getLastUpdated() { + if (lastUpdatedBuilder_ == null) { + return lastUpdated_ == null ? com.google.protobuf.Timestamp.getDefaultInstance() : lastUpdated_; + } else { + return lastUpdatedBuilder_.getMessage(); + } + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + public Builder setLastUpdated(com.google.protobuf.Timestamp value) { + if (lastUpdatedBuilder_ == null) { + if (value == null) { + throw new NullPointerException(); + } + lastUpdated_ = value; + } else { + lastUpdatedBuilder_.setMessage(value); + } + bitField0_ |= 0x00000080; + onChanged(); + return this; + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + public Builder setLastUpdated( + com.google.protobuf.Timestamp.Builder builderForValue) { + if (lastUpdatedBuilder_ == null) { + lastUpdated_ = builderForValue.build(); + } else { + lastUpdatedBuilder_.setMessage(builderForValue.build()); + } + bitField0_ |= 0x00000080; + onChanged(); + return this; + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + public Builder mergeLastUpdated(com.google.protobuf.Timestamp value) { + if (lastUpdatedBuilder_ == null) { + if (((bitField0_ & 0x00000080) != 0) && + lastUpdated_ != null && + lastUpdated_ != com.google.protobuf.Timestamp.getDefaultInstance()) { + getLastUpdatedBuilder().mergeFrom(value); + } else { + lastUpdated_ = value; + } + } else { + lastUpdatedBuilder_.mergeFrom(value); + } + bitField0_ |= 0x00000080; + onChanged(); + return this; + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + public Builder clearLastUpdated() { + bitField0_ = (bitField0_ & ~0x00000080); + lastUpdated_ = null; + if (lastUpdatedBuilder_ != null) { + lastUpdatedBuilder_.dispose(); + lastUpdatedBuilder_ = null; + } + onChanged(); + return this; + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + public com.google.protobuf.Timestamp.Builder getLastUpdatedBuilder() { + bitField0_ |= 0x00000080; + onChanged(); + return getLastUpdatedFieldBuilder().getBuilder(); + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + public com.google.protobuf.TimestampOrBuilder getLastUpdatedOrBuilder() { + if (lastUpdatedBuilder_ != null) { + return lastUpdatedBuilder_.getMessageOrBuilder(); + } else { + return lastUpdated_ == null ? + com.google.protobuf.Timestamp.getDefaultInstance() : lastUpdated_; + } + } + /** + * .google.protobuf.Timestamp last_updated = 8; + */ + private com.google.protobuf.SingleFieldBuilderV3< + com.google.protobuf.Timestamp, com.google.protobuf.Timestamp.Builder, com.google.protobuf.TimestampOrBuilder> + getLastUpdatedFieldBuilder() { + if (lastUpdatedBuilder_ == null) { + lastUpdatedBuilder_ = new com.google.protobuf.SingleFieldBuilderV3< + com.google.protobuf.Timestamp, com.google.protobuf.Timestamp.Builder, com.google.protobuf.TimestampOrBuilder>( + getLastUpdated(), + getParentForChildren(), + isClean()); + lastUpdated_ = null; + } + return lastUpdatedBuilder_; + } + @java.lang.Override + public final Builder setUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.setUnknownFields(unknownFields); + } + + @java.lang.Override + public final Builder mergeUnknownFields( + final com.google.protobuf.UnknownFieldSet unknownFields) { + return super.mergeUnknownFields(unknownFields); + } + + + // @@protoc_insertion_point(builder_scope:kafka.automq.table.process.proto.Person) + } + + // @@protoc_insertion_point(class_scope:kafka.automq.table.process.proto.Person) + private static final kafka.automq.table.process.proto.PersonProto.Person DEFAULT_INSTANCE; + static { + DEFAULT_INSTANCE = new kafka.automq.table.process.proto.PersonProto.Person(); + } + + public static kafka.automq.table.process.proto.PersonProto.Person getDefaultInstance() { + return DEFAULT_INSTANCE; + } + + private static final com.google.protobuf.Parser + PARSER = new com.google.protobuf.AbstractParser() { + @java.lang.Override + public Person parsePartialFrom( + com.google.protobuf.CodedInputStream input, + com.google.protobuf.ExtensionRegistryLite extensionRegistry) + throws com.google.protobuf.InvalidProtocolBufferException { + Builder builder = newBuilder(); + try { + builder.mergeFrom(input, extensionRegistry); + } catch (com.google.protobuf.InvalidProtocolBufferException e) { + throw e.setUnfinishedMessage(builder.buildPartial()); + } catch (com.google.protobuf.UninitializedMessageException e) { + throw e.asInvalidProtocolBufferException().setUnfinishedMessage(builder.buildPartial()); + } catch (java.io.IOException e) { + throw new com.google.protobuf.InvalidProtocolBufferException(e) + .setUnfinishedMessage(builder.buildPartial()); + } + return builder.buildPartial(); + } + }; + + public static com.google.protobuf.Parser parser() { + return PARSER; + } + + @java.lang.Override + public com.google.protobuf.Parser getParserForType() { + return PARSER; + } + + @java.lang.Override + public kafka.automq.table.process.proto.PersonProto.Person getDefaultInstanceForType() { + return DEFAULT_INSTANCE; + } + + } + + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_kafka_automq_table_process_proto_Address_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_kafka_automq_table_process_proto_Address_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_kafka_automq_table_process_proto_Person_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_kafka_automq_table_process_proto_Person_fieldAccessorTable; + private static final com.google.protobuf.Descriptors.Descriptor + internal_static_kafka_automq_table_process_proto_Person_AttributesEntry_descriptor; + private static final + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable + internal_static_kafka_automq_table_process_proto_Person_AttributesEntry_fieldAccessorTable; + + public static com.google.protobuf.Descriptors.FileDescriptor + getDescriptor() { + return descriptor; + } + private static com.google.protobuf.Descriptors.FileDescriptor + descriptor; + static { + java.lang.String[] descriptorData = { + "\n\022proto/person.proto\022 kafka.automq.table" + + ".process.proto\032\037google/protobuf/timestam" + + "p.proto\"\'\n\007Address\022\016\n\006street\030\001 \001(\t\022\014\n\004ci" + + "ty\030\002 \001(\t\"\307\002\n\006Person\022\n\n\002id\030\001 \001(\003\022\014\n\004name\030" + + "\002 \001(\t\022\021\n\tis_active\030\003 \001(\010\022\022\n\nextra_data\030\004" + + " \001(\014\022:\n\007address\030\005 \001(\0132).kafka.automq.tab" + + "le.process.proto.Address\022\r\n\005roles\030\006 \003(\t\022" + + "L\n\nattributes\030\007 \003(\01328.kafka.automq.table" + + ".process.proto.Person.AttributesEntry\0220\n" + + "\014last_updated\030\010 \001(\0132\032.google.protobuf.Ti" + + "mestamp\0321\n\017AttributesEntry\022\013\n\003key\030\001 \001(\t\022" + + "\r\n\005value\030\002 \001(\t:\0028\001B/\n kafka.automq.table" + + ".process.protoB\013PersonProtob\006proto3" + }; + descriptor = com.google.protobuf.Descriptors.FileDescriptor + .internalBuildGeneratedFileFrom(descriptorData, + new com.google.protobuf.Descriptors.FileDescriptor[] { + com.google.protobuf.TimestampProto.getDescriptor(), + }); + internal_static_kafka_automq_table_process_proto_Address_descriptor = + getDescriptor().getMessageTypes().get(0); + internal_static_kafka_automq_table_process_proto_Address_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_kafka_automq_table_process_proto_Address_descriptor, + new java.lang.String[] { "Street", "City", }); + internal_static_kafka_automq_table_process_proto_Person_descriptor = + getDescriptor().getMessageTypes().get(1); + internal_static_kafka_automq_table_process_proto_Person_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_kafka_automq_table_process_proto_Person_descriptor, + new java.lang.String[] { "Id", "Name", "IsActive", "ExtraData", "Address", "Roles", "Attributes", "LastUpdated", }); + internal_static_kafka_automq_table_process_proto_Person_AttributesEntry_descriptor = + internal_static_kafka_automq_table_process_proto_Person_descriptor.getNestedTypes().get(0); + internal_static_kafka_automq_table_process_proto_Person_AttributesEntry_fieldAccessorTable = new + com.google.protobuf.GeneratedMessageV3.FieldAccessorTable( + internal_static_kafka_automq_table_process_proto_Person_AttributesEntry_descriptor, + new java.lang.String[] { "Key", "Value", }); + com.google.protobuf.TimestampProto.getDescriptor(); + } + + // @@protoc_insertion_point(outer_class_scope) +} diff --git a/core/src/test/java/kafka/automq/table/process/transform/DebeziumUnwrapTransformTest.java b/core/src/test/java/kafka/automq/table/process/transform/DebeziumUnwrapTransformTest.java new file mode 100644 index 0000000000..b673988f93 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/transform/DebeziumUnwrapTransformTest.java @@ -0,0 +1,245 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.transform; + +import kafka.automq.table.process.TransformContext; +import kafka.automq.table.process.exception.TransformException; + +import org.apache.kafka.common.record.Record; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +class DebeziumUnwrapTransformTest { + + private static final Schema ROW_SCHEMA = + SchemaBuilder.record("row") + .fields() + .requiredLong("account_id") + .requiredDouble("balance") + .requiredString("last_updated") + .endRecord(); + + private static final Schema SOURCE_SCHEMA = + SchemaBuilder.record("source") + .fields() + .requiredString("db") + .optionalString("schema") + .requiredString("table") + .endRecord(); + + private static final Schema DEBEZIUM_SCHEMA = + SchemaBuilder.record("debezium_event") + .fields() + .requiredString("op") + .optionalLong("ts_ms") + .name("source").type(SOURCE_SCHEMA).noDefault() + .name("before").type().unionOf().nullType().and().type(ROW_SCHEMA).endUnion().nullDefault() + .name("after").type().unionOf().nullType().and().type(ROW_SCHEMA).endUnion().nullDefault() + .endRecord(); + + private static final Schema SOURCE_SCHEMA_WITHOUT_SCHEMA_FIELD = + SchemaBuilder.record("source_no_schema") + .fields() + .requiredString("db") + .requiredString("table") + .endRecord(); + + private static final Schema DEBEZIUM_SCHEMA_WITHOUT_SCHEMA_IN_SOURCE = + SchemaBuilder.record("debezium_event_no_source_schema") + .fields() + .requiredString("op") + .optionalLong("ts_ms") + .name("source").type(SOURCE_SCHEMA_WITHOUT_SCHEMA_FIELD).noDefault() + .name("before").type().unionOf().nullType().and().type(ROW_SCHEMA).endUnion().nullDefault() + .name("after").type().unionOf().nullType().and().type(ROW_SCHEMA).endUnion().nullDefault() + .endRecord(); + + private DebeziumUnwrapTransform transform; + private FlattenTransform flattenTransform; + private TransformContext context; + private Record kafkaRecord; + + @BeforeEach + void setUp() { + transform = new DebeziumUnwrapTransform(); + transform.configure(Collections.emptyMap()); + + flattenTransform = new FlattenTransform(); + flattenTransform.configure(Collections.emptyMap()); + + context = mock(TransformContext.class); + kafkaRecord = mock(Record.class); + when(context.getKafkaRecord()).thenReturn(kafkaRecord); + when(kafkaRecord.offset()).thenReturn(123L); + } + + @Test + void testNullRecord() { + assertThrows(NullPointerException.class, () -> transform.apply(null, context)); + } + + @Test + void testCreateOperation() throws TransformException { + GenericRecord event = createDebeziumEvent("c", 1L, 100.0); + GenericRecord result = transform.apply(event, context); + + assertNotNull(result); + assertEquals(1L, result.get("account_id")); + assertEquals(100.0, result.get("balance")); + GenericRecord cdc = (GenericRecord) result.get("_cdc"); + assertNotNull(cdc); + assertEquals("I", cdc.get("op")); + assertEquals(123L, cdc.get("offset")); + assertTrue(cdc.get("ts") instanceof Long); + assertEquals("test_schema.test_table", cdc.get("source")); + } + + @Test + void testUpdateOperation() throws TransformException { + GenericRecord event = createDebeziumEvent("u", 2L, 200.0); + GenericRecord result = transform.apply(event, context); + + assertNotNull(result); + assertEquals(2L, result.get("account_id")); + assertEquals(200.0, result.get("balance")); + GenericRecord cdc = (GenericRecord) result.get("_cdc"); + assertNotNull(cdc); + assertEquals("U", cdc.get("op")); + } + + @Test + void testDeleteOperation() throws TransformException { + GenericRecord event = createDebeziumEvent("d", 3L, 300.0); + GenericRecord result = transform.apply(event, context); + + assertNotNull(result); + assertEquals(3L, result.get("account_id")); + assertEquals(300.0, result.get("balance")); + GenericRecord cdc = (GenericRecord) result.get("_cdc"); + assertNotNull(cdc); + assertEquals("D", cdc.get("op")); + } + + @Test + void testDeleteWithNullBefore() { + GenericRecordBuilder builder = new GenericRecordBuilder(DEBEZIUM_SCHEMA) + .set("op", "d") + .set("ts_ms", System.currentTimeMillis()) + .set("source", createSourceRecord()) + .set("before", null) + .set("after", null); + GenericRecord event = builder.build(); + + TransformException e = assertThrows(TransformException.class, () -> transform.apply(event, context)); + assertTrue(e.getMessage().contains("Invalid DELETE record: missing required 'before' data")); + } + + + @Test + void testSourceWithNullSchema() throws TransformException { + GenericRecord sourceWithNullSchema = new GenericRecordBuilder(SOURCE_SCHEMA) + .set("db", "test_db_from_db_field") + .set("schema", null) + .set("table", "test_table") + .build(); + + GenericRecord event = createDebeziumEvent("c", 1L, 100.0, sourceWithNullSchema); + GenericRecord result = transform.apply(event, context); + + GenericRecord cdc = (GenericRecord) result.get("_cdc"); + assertNotNull(cdc); + assertEquals("test_db_from_db_field.test_table", cdc.get("source")); + } + + @Test + void testSourceWithoutSchemaField() throws TransformException { + GenericRecord sourceWithoutSchemaField = new GenericRecordBuilder(SOURCE_SCHEMA_WITHOUT_SCHEMA_FIELD) + .set("db", "db_only") + .set("table", "table_from_db_only") + .build(); + + GenericRecord row = createRowRecord(1L, 100.0); + GenericRecord event = new GenericRecordBuilder(DEBEZIUM_SCHEMA_WITHOUT_SCHEMA_IN_SOURCE) + .set("op", "c") + .set("ts_ms", System.currentTimeMillis()) + .set("source", sourceWithoutSchemaField) + .set("after", row) + .set("before", null) + .build(); + + GenericRecord result = transform.apply(event, context); + + GenericRecord cdc = (GenericRecord) result.get("_cdc"); + assertNotNull(cdc); + assertEquals("db_only.table_from_db_only", cdc.get("source")); + } + + private GenericRecord createDebeziumEvent(String op, long accountId, double balance) { + return createDebeziumEvent(op, accountId, balance, createSourceRecord()); + } + + private GenericRecord createDebeziumEvent(String op, long accountId, double balance, GenericRecord source) { + GenericRecord row = createRowRecord(accountId, balance); + GenericRecordBuilder builder = new GenericRecordBuilder(DEBEZIUM_SCHEMA) + .set("op", op) + .set("ts_ms", System.currentTimeMillis()) + .set("source", source); + + if ("c".equals(op) || "r".equals(op) || "u".equals(op)) { + builder.set("after", row); + } + if ("d".equals(op) || "u".equals(op)) { + builder.set("before", row); + } + return builder.build(); + } + + private GenericRecord createRowRecord(long accountId, double balance) { + return new GenericRecordBuilder(ROW_SCHEMA) + .set("account_id", accountId) + .set("balance", balance) + .set("last_updated", "2025-01-01T00:00:00Z") + .build(); + } + + private GenericRecord createSourceRecord() { + return new GenericRecordBuilder(SOURCE_SCHEMA) + .set("db", "test_db") + .set("schema", "test_schema") + .set("table", "test_table") + .build(); + } +} diff --git a/core/src/test/java/kafka/automq/table/process/transform/FlattenTransformTest.java b/core/src/test/java/kafka/automq/table/process/transform/FlattenTransformTest.java new file mode 100644 index 0000000000..7677530e6f --- /dev/null +++ b/core/src/test/java/kafka/automq/table/process/transform/FlattenTransformTest.java @@ -0,0 +1,107 @@ + +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.process.transform; + +import kafka.automq.table.process.RecordAssembler; +import kafka.automq.table.process.TransformContext; +import kafka.automq.table.process.exception.TransformException; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +@Tag("S3Unit") +class FlattenTransformTest { + + private static final Schema INNER_SCHEMA = + SchemaBuilder.record("Inner") + .fields() + .requiredLong("id") + .requiredString("data") + .endRecord(); + + private static final Schema VALUE_CONTAINER_SCHEMA = + SchemaBuilder.record("ValueContainer") + .fields() + .name(RecordAssembler.KAFKA_VALUE_FIELD).type(INNER_SCHEMA).noDefault() + .endRecord(); + + private static final Schema NON_WRAPPED_SCHEMA = + SchemaBuilder.record("NonWrapped") + .fields() + .requiredLong("id") + .endRecord(); + + private static final Schema STRING_VALUE_SCHEMA = + SchemaBuilder.record("StringValue") + .fields() + .requiredString(RecordAssembler.KAFKA_VALUE_FIELD) + .endRecord(); + + + private FlattenTransform transform; + private TransformContext context; + + @BeforeEach + void setUp() { + transform = new FlattenTransform(); + transform.configure(Collections.emptyMap()); + context = mock(TransformContext.class); + } + + @Test + void testApplyWhenRecordIsWrappedShouldUnwrap() throws TransformException { + GenericRecord innerRecord = new GenericRecordBuilder(INNER_SCHEMA) + .set("id", 1L) + .set("data", "test") + .build(); + + GenericRecord outerRecord = new GenericRecordBuilder(VALUE_CONTAINER_SCHEMA) + .set(RecordAssembler.KAFKA_VALUE_FIELD, innerRecord) + .build(); + + GenericRecord result = transform.apply(outerRecord, context); + + assertSame(innerRecord, result, "The transform should return the inner record."); + } + + @Test + void testApplyWhenRecordIsNotWrappedShouldThrowException() { + GenericRecord nonWrappedRecord = new GenericRecordBuilder(NON_WRAPPED_SCHEMA) + .set("id", 123L) + .build(); + + TransformException e = assertThrows(TransformException.class, () -> transform.apply(nonWrappedRecord, context)); + assertTrue(e.getMessage().contains("Record is null or has no value field")); + } + +} diff --git a/core/src/test/java/kafka/automq/table/utils/PartitionUtilTest.java b/core/src/test/java/kafka/automq/table/utils/PartitionUtilTest.java new file mode 100644 index 0000000000..0e8223497e --- /dev/null +++ b/core/src/test/java/kafka/automq/table/utils/PartitionUtilTest.java @@ -0,0 +1,65 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.utils; + +import com.google.common.collect.ImmutableMap; + +import org.apache.iceberg.PartitionField; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.util.List; + + +@Tag("S3Unit") +public class PartitionUtilTest { + + @Test + public void testEvolve() { + InMemoryCatalog inMemoryCatalog = new InMemoryCatalog(); + inMemoryCatalog.initialize("test", ImmutableMap.of()); + inMemoryCatalog.createNamespace(Namespace.of("default")); + Schema v1Schema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "timestamp", Types.TimestampType.withZone()) + ); + Table table = inMemoryCatalog.createTable(TableIdentifier.parse("default.test"), v1Schema); + // add partition + PartitionUtil.evolve(List.of("id", "bucket(name, 8)", "hour(timestamp)"), table); + Assertions.assertEquals(List.of("id", "name_bucket_8", "timestamp_hour"), table.spec().fields().stream().map(PartitionField::name).toList()); + + // replace partition + PartitionUtil.evolve(List.of("id", "bucket(name, 8)", "day(timestamp)"), table); + Assertions.assertEquals(List.of("id", "name_bucket_8", "timestamp_day"), table.spec().fields().stream().map(PartitionField::name).toList()); + + // drop partition + PartitionUtil.evolve(List.of("bucket(name, 8)", "day(timestamp)"), table); + Assertions.assertEquals(List.of("name_bucket_8", "timestamp_day"), table.spec().fields().stream().map(PartitionField::name).toList()); + } + +} diff --git a/core/src/test/java/kafka/automq/table/worker/IcebergSchemaChangeCollectorTest.java b/core/src/test/java/kafka/automq/table/worker/IcebergSchemaChangeCollectorTest.java new file mode 100644 index 0000000000..7e63451603 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/IcebergSchemaChangeCollectorTest.java @@ -0,0 +1,621 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import com.google.common.collect.ImmutableMap; + +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +public class IcebergSchemaChangeCollectorTest { + private InMemoryCatalog catalog; + + @BeforeEach + public void setup() { + catalog = initializeCatalog(); + catalog.createNamespace(Namespace.of("default")); + } + + @Test + public void shouldReturnEmptyWhenSchemasMatch() { + Schema schema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + List changes = collectChanges(schema, schema); + assertTrue(changes.isEmpty()); + } + + @Test + public void shouldDetectTopLevelAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "name", null); + assertEquals(Types.StringType.get(), change.getNewType()); + } + + @Test + public void shouldDetectTopLevelOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.optional(1, "id", Types.IntegerType.get())); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "id", null); + } + + @Test + public void shouldDetectTopLevelPromotion() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.LongType.get())); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE, "id", null); + assertEquals(Types.LongType.get(), change.getNewType()); + } + + @Test + public void shouldDetectNestedAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get())))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get()), + Types.NestedField.optional(4, "email", Types.StringType.get())))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "email", "user"); + } + + @Test + public void shouldDetectListElementStructAddition() { + Schema initial = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "street", Types.StringType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of( + Types.NestedField.optional(3, "street", Types.StringType.get()), + Types.NestedField.optional(4, "zip", Types.IntegerType.get()))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "zip", "addresses.element"); + } + + @Test + public void shouldPromoteListElementStructFieldType() { + Schema initial = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.IntegerType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.LongType.get()))))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE, "zip", "addresses.element"); + assertEquals(Types.LongType.get(), change.getNewType()); + } + + @Test + public void shouldPromoteMapValueStructFieldType() { + Schema initial = new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of(Types.NestedField.optional(4, "zip", Types.IntegerType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of(Types.NestedField.optional(4, "zip", Types.LongType.get()))))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE, "zip", "attributes.value"); + assertEquals(Types.LongType.get(), change.getNewType()); + } + + @Test + public void shouldMakeListElementStructFieldOptional() { + Schema initial = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.required(3, "zip", Types.IntegerType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.IntegerType.get()))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "zip", "addresses.element"); + } + + @Test + public void shouldSoftRemoveMapValueStructField() { + Schema initial = new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of(Types.NestedField.required(4, "zip", Types.IntegerType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of()))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "zip", "attributes.value"); + } + + @Test + public void shouldSkipDuplicateListElementChanges() { + Schema schema = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.LongType.get()))))); + + List changes = List.of( + new IcebergTableManager.SchemaChange(IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "zip", null, "addresses.element"), + new IcebergTableManager.SchemaChange(IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE, "zip", Types.LongType.get(), "addresses.element")); + + IcebergTableManager manager = new IcebergTableManager(initializeCatalog(), TableIdentifier.of("default", "dummy"), mock(WorkerConfig.class)); + Table table = manager.getTableOrCreate(schema); + manager.applySchemaChange(table, changes); + } + + @Test + public void ignoresElementTypeReplacementAsIncompatible() { + Schema initial = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.IntegerType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, Types.IntegerType.get()))); + + List changes = collectChanges(initial, current); + assertTrue(changes.isEmpty(), "Incompatible list element type change should be ignored"); + } + + @Test + public void ignoresPrimitiveListElementPromotion() { + Schema initial = new Schema( + Types.NestedField.optional(1, "nums", Types.ListType.ofOptional(2, Types.IntegerType.get()))); + Schema current = new Schema( + Types.NestedField.optional(1, "nums", Types.ListType.ofOptional(2, Types.LongType.get()))); + + List changes = collectChanges(initial, current); + assertTrue(changes.isEmpty(), "list -> list promotion is unsupported and should be ignored"); + } + + @Test + public void shouldDetectMapValueStructAddition() { + Schema initial = new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of(Types.NestedField.optional(4, "city", Types.StringType.get()))))); + Schema current = new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of( + Types.NestedField.optional(4, "city", Types.StringType.get()), + Types.NestedField.optional(5, "country", Types.StringType.get()))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "country", "attributes.value"); + } + + @Test + public void shouldDetectDeepNestedAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "name", Types.StringType.get())))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "name", Types.StringType.get()), + Types.NestedField.optional(5, "age", Types.IntegerType.get())))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "age", "user.profile"); + } + + @Test + public void shouldDetectNestedFieldMadeOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "name", Types.StringType.get()), + Types.NestedField.required(5, "age", Types.IntegerType.get())))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "name", Types.StringType.get()), + Types.NestedField.optional(5, "age", Types.IntegerType.get())))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "age", "user.profile"); + } + + @Test + public void shouldDetectNestedPromotion() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "age", Types.IntegerType.get())))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "age", Types.LongType.get())))))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE, "age", "user.profile"); + assertEquals(Types.LongType.get(), change.getNewType()); + } + + @Test + public void shouldSoftRemoveTopLevelField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "email", Types.StringType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "email", null); + } + + @Test + public void shouldSoftRemoveNonPrimitiveField() { + Schema initial = new Schema( + Types.NestedField.required(1, "ids", Types.ListType.ofRequired(2, Types.LongType.get()))); + Schema current = new Schema(); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "ids", null); + // ensure newType is null for non-primitive types so runtime won't call asPrimitiveType + assertEquals(null, change.getNewType()); + } + + @Test + public void shouldSoftRemoveNestedField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "name", Types.StringType.get()), + Types.NestedField.required(5, "age", Types.IntegerType.get())))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "profile", Types.StructType.of( + Types.NestedField.required(4, "name", Types.StringType.get())))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "age", "user.profile"); + } + + @Test + public void shouldReportMixedChanges() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "score", Types.FloatType.get()), + Types.NestedField.optional(4, "old_field", Types.StringType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "score", Types.DoubleType.get()), + Types.NestedField.optional(5, "new_field", Types.StringType.get())); + + List changes = collectChanges(initial, current); + assertEquals(4, changes.size()); + assertTrue(changes.stream().anyMatch(c -> c.getType() == IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE + && c.getColumnName().equals("id"))); + assertTrue(changes.stream().anyMatch(c -> c.getType() == IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL + && c.getColumnName().equals("name"))); + assertTrue(changes.stream().anyMatch(c -> c.getType() == IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE + && c.getColumnName().equals("score"))); + assertTrue(changes.stream().anyMatch(c -> c.getType() == IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN + && c.getColumnName().equals("new_field"))); + } + + + @Test + public void shouldDetectOptionalListFieldAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "tags", Types.ListType.ofOptional(3, Types.StringType.get()))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "tags", null); + assertEquals(Types.ListType.ofOptional(3, Types.StringType.get()), change.getNewType()); + } + + @Test + public void shouldDetectOptionalMapFieldAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "metadata", Types.MapType.ofOptional(3, 4, + Types.StringType.get(), Types.StringType.get()))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "metadata", null); + assertEquals(Types.MapType.ofOptional(3, 4, Types.StringType.get(), Types.StringType.get()), + change.getNewType()); + } + + @Test + public void shouldDetectRequiredListFieldAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "scores", Types.ListType.ofRequired(3, Types.IntegerType.get()))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "scores", null); + assertEquals(Types.ListType.ofRequired(3, Types.IntegerType.get()), change.getNewType()); + } + + @Test + public void shouldDetectRequiredMapFieldAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "attributes", Types.MapType.ofRequired(3, 4, + Types.StringType.get(), Types.IntegerType.get()))); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "attributes", null); + assertEquals(Types.MapType.ofRequired(3, 4, Types.StringType.get(), Types.IntegerType.get()), + change.getNewType()); + } + + @Test + public void shouldSoftRemoveRequiredListField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "tags", Types.ListType.ofOptional(3, Types.StringType.get()))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "tags", null); + // newType should be null for non-primitive types + assertEquals(null, change.getNewType()); + } + + @Test + public void shouldSoftRemoveRequiredMapField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "metadata", Types.MapType.ofOptional(3, 4, + Types.StringType.get(), Types.StringType.get()))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + IcebergTableManager.SchemaChange change = assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "metadata", null); + assertEquals(null, change.getNewType()); + } + + @Test + public void shouldSoftRemoveOptionalListField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "tags", Types.ListType.ofOptional(3, Types.StringType.get()))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + // Optional field removal should still trigger MAKE_OPTIONAL to ensure idempotency + List changes = collectChanges(initial, current); + assertTrue(changes.isEmpty() || changes.stream().allMatch(c -> + c.getType() == IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL && + c.getColumnName().equals("tags"))); + } + + @Test + public void shouldSoftRemoveOptionalMapField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "metadata", Types.MapType.ofOptional(3, 4, + Types.StringType.get(), Types.StringType.get()))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + List changes = collectChanges(initial, current); + assertTrue(changes.isEmpty() || changes.stream().allMatch(c -> + c.getType() == IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL && + c.getColumnName().equals("metadata"))); + } + + @Test + public void shouldMakeListFieldOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "tags", Types.ListType.ofOptional(3, Types.StringType.get()))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "tags", Types.ListType.ofOptional(3, Types.StringType.get()))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "tags", null); + } + + @Test + public void shouldMakeMapFieldOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "metadata", Types.MapType.ofOptional(3, 4, + Types.StringType.get(), Types.StringType.get()))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "metadata", Types.MapType.ofOptional(3, 4, + Types.StringType.get(), Types.StringType.get()))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "metadata", null); + } + + @Test + public void shouldDetectNestedListFieldAddition() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get())))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get()), + Types.NestedField.optional(4, "hobbies", Types.ListType.ofOptional(5, Types.StringType.get()))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, "hobbies", "user"); + } + + @Test + public void shouldSoftRemoveNestedListField() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get()), + Types.NestedField.required(4, "hobbies", Types.ListType.ofOptional(5, Types.StringType.get()))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get())))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "hobbies", "user"); + } + + @Test + public void shouldMakeNestedMapFieldOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get()), + Types.NestedField.required(4, "preferences", Types.MapType.ofOptional(5, 6, + Types.StringType.get(), Types.StringType.get()))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "user", Types.StructType.of( + Types.NestedField.required(3, "name", Types.StringType.get()), + Types.NestedField.optional(4, "preferences", Types.MapType.ofOptional(5, 6, + Types.StringType.get(), Types.StringType.get()))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "preferences", "user"); + } + + + @Test + public void shouldMakeStructFieldOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "address", Types.StructType.of( + Types.NestedField.required(3, "street", Types.StringType.get()), + Types.NestedField.required(4, "city", Types.StringType.get())))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "address", Types.StructType.of( + Types.NestedField.required(3, "street", Types.StringType.get()), + Types.NestedField.required(4, "city", Types.StringType.get())))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "address", null); + } + + @Test + public void shouldMakeListNestedStructFieldOptional() { + Schema initial = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "contacts", Types.ListType.ofRequired(3, + Types.StructType.of( + Types.NestedField.required(4, "type", Types.StringType.get()), + Types.NestedField.required(5, "detail", Types.StringType.get()))))); + Schema current = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.required(2, "contacts", Types.ListType.ofRequired(3, + Types.StructType.of( + Types.NestedField.optional(4, "type", Types.StringType.get()), + Types.NestedField.required(5, "detail", Types.StringType.get()))))); + + assertSingleChange( + collectChanges(initial, current), IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, "type", "contacts.element"); + } + + private List collectChanges(Schema tableSchema, Schema currentSchema) { + TableIdentifier tableId = TableIdentifier.of("default", generateRandomTableName()); + IcebergTableManager tableManager = new IcebergTableManager(catalog, tableId, mock(WorkerConfig.class)); + Table table = tableManager.getTableOrCreate(tableSchema); + return tableManager.collectSchemaChanges(currentSchema, table); + } + + private IcebergTableManager.SchemaChange assertSingleChange(List changes, + IcebergTableManager.SchemaChange.ChangeType type, + String columnName, + String parentName) { + assertEquals(1, changes.size(), "Expected exactly one schema change"); + IcebergTableManager.SchemaChange change = changes.get(0); + assertEquals(type, change.getType()); + assertEquals(columnName, change.getColumnName()); + assertEquals(parentName, change.getParentName()); + return change; + } + + private String generateRandomTableName() { + int randomNum = ThreadLocalRandom.current().nextInt(1000, 10000); + return "schema_table_" + randomNum; + } + + private InMemoryCatalog initializeCatalog() { + InMemoryCatalog inMemoryCatalog = new InMemoryCatalog(); + inMemoryCatalog.initialize("test", ImmutableMap.of()); + return inMemoryCatalog; + } +} diff --git a/core/src/test/java/kafka/automq/table/worker/IcebergTableManagerTest.java b/core/src/test/java/kafka/automq/table/worker/IcebergTableManagerTest.java new file mode 100644 index 0000000000..f52a5b082c --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/IcebergTableManagerTest.java @@ -0,0 +1,498 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import com.google.common.collect.ImmutableMap; + +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.UpdateSchema; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.exceptions.NoSuchTableException; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; + +import java.io.IOException; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +public class IcebergTableManagerTest { + private InMemoryCatalog catalog; + + @BeforeEach + public void setup() { + catalog = initializeCatalog(); + catalog.createNamespace(Namespace.of("default")); + } + + @Test + public void shouldCreateTableOnceAndReuseInstance() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Schema schema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + Table table = manager.getTableOrCreate(schema); + assertNotNull(table); + assertEquals(table, manager.getTableOrCreate(schema)); + } + + @Test + public void createsMissingNamespaceBeforeCreatingTable() { + String namespace = "ns_" + ThreadLocalRandom.current().nextInt(1000, 10000); + TableIdentifier tableId = TableIdentifier.of(namespace, "table_" + ThreadLocalRandom.current().nextInt(1000, 10000)); + IcebergTableManager manager = newManager(tableId); + + Schema schema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + + Table table = manager.getTableOrCreate(schema); + assertNotNull(table); + // Namespace should now exist and table should be loadable + assertNotNull(catalog.loadTable(tableId)); + } + + @Test + public void supportsFieldNamesContainingDots() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Schema schema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name.test", Types.StringType.get())); + + Table table = manager.getTableOrCreate(schema); + Types.NestedField field = table.schema().findField("name.test"); + assertNotNull(field); + assertEquals("name.test", field.name()); + } + + @Test + public void addsPrimitiveColumn() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "email", Types.StringType.get()))); + + assertNotNull(updated.schema().findField("email")); + } + + @Test + public void promotesPrimitiveColumnType() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()))); + + assertEquals(Types.LongType.get(), updated.schema().findField("id").type()); + } + + @Test + public void makesColumnOptional() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.optional(1, "id", Types.IntegerType.get()))); + + assertTrue(updated.schema().findField("id").isOptional()); + } + + @Test + public void addsStructColumn() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "address", Types.StructType.of( + Types.NestedField.optional(3, "street", Types.StringType.get()), + Types.NestedField.optional(4, "zipCode", Types.IntegerType.get()))))); + + Types.NestedField address = updated.schema().findField("address"); + assertNotNull(address); + assertTrue(address.type().isStructType()); + assertNotNull(updated.schema().findField("address.street")); + assertNotNull(updated.schema().findField("address.zipCode")); + } + + @Test + public void addsNestedFieldInsideStruct() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "address", Types.StructType.of( + Types.NestedField.optional(3, "street", Types.StringType.get()))))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "address", Types.StructType.of( + Types.NestedField.optional(3, "street", Types.StringType.get()), + Types.NestedField.optional(4, "zipCode", Types.IntegerType.get()))))); + + assertNotNull(updated.schema().findField("address.zipCode")); + } + + @Test + public void addsFieldInsideListElementStruct() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "street", Types.StringType.get())))))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of( + Types.NestedField.optional(3, "street", Types.StringType.get()), + Types.NestedField.optional(4, "zip", Types.IntegerType.get())))))); + + assertNotNull(updated.schema().findField("addresses.element.zip")); + } + + @Test + public void addsFieldInsideMapValueStruct() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of(Types.NestedField.optional(4, "city", Types.StringType.get())))))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.optional(1, "attributes", Types.MapType.ofOptional(2, 3, + Types.StringType.get(), + Types.StructType.of( + Types.NestedField.optional(4, "city", Types.StringType.get()), + Types.NestedField.optional(5, "country", Types.StringType.get())))))); + + assertNotNull(updated.schema().findField("attributes.value.country")); + } + + @Test + public void promotesFieldInsideListElementStruct() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.IntegerType.get())))))); + + Table updated = applyChanges(manager, table, new Schema( + Types.NestedField.optional(1, "addresses", Types.ListType.ofOptional(2, + Types.StructType.of(Types.NestedField.optional(3, "zip", Types.LongType.get())))))); + + assertEquals(Types.LongType.get(), updated.schema().findField("addresses.element.zip").type()); + } + + @Test + public void exposesPartitionSpecAndAllowsReset() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + assertNotNull(manager.spec(), "Partition spec should be captured after table creation"); + + manager.reset(); + Table reloaded = manager.getTableOrCreate(table.schema()); + assertEquals(table.schema().asStruct(), reloaded.schema().asStruct()); + } + + @Test + public void handleSchemaChangesWithFlushTriggersFlushOnlyWhenNeeded() throws Exception { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + Schema newSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + AtomicInteger flushCount = new AtomicInteger(); + boolean changed = manager.handleSchemaChangesWithFlush(newSchema, () -> { + flushCount.incrementAndGet(); + }); + assertTrue(changed); + assertEquals(1, flushCount.get()); + assertNotNull(catalog.loadTable(tableId).schema().findField("name")); + + boolean noChange = manager.handleSchemaChangesWithFlush(newSchema, () -> { + flushCount.incrementAndGet(); + }); + assertFalse(noChange); + assertEquals(1, flushCount.get()); + } + + @Test + public void handleSchemaChangesWithFlushPropagatesFlushFailures() throws Exception { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()))); + + Schema newSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + assertThrows(IOException.class, () -> manager.handleSchemaChangesWithFlush(newSchema, () -> { + throw new IOException("flush failed"); + })); + + assertNull(catalog.loadTable(tableId).schema().findField("name")); + } + + @Test + public void retriesSchemaCommitOnFailure() { + TableIdentifier tableId = randomTableId(); + Schema baseSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + Schema updatedSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "email", Types.StringType.get())); + + Catalog mockCatalog = mock(Catalog.class); + Table mockTable = mock(Table.class); + UpdateSchema mockUpdateSchema = mock(UpdateSchema.class); + + when(mockCatalog.loadTable(eq(tableId))).thenThrow(new NoSuchTableException("Table not found")); + when(mockCatalog.createTable(eq(tableId), eq(baseSchema), any(), any())).thenReturn(mockTable); + when(mockTable.schema()).thenReturn(baseSchema); + when(mockTable.updateSchema()).thenReturn(mockUpdateSchema); + when(mockUpdateSchema.addColumn("email", Types.StringType.get())).thenReturn(mockUpdateSchema); + + doAnswer(new Answer() { + private int count = 0; + + @Override + public Void answer(InvocationOnMock invocation) { + if (count < 1) { + count++; + throw new RuntimeException("Commit Error"); + } + return null; + } + }).when(mockUpdateSchema).commit(); + + IcebergTableManager manager = new IcebergTableManager(mockCatalog, tableId, mock(WorkerConfig.class)); + Table table = manager.getTableOrCreate(baseSchema); + + Record record = GenericRecord.create(updatedSchema); + List schemaChanges = manager.checkSchemaChanges(table, record.struct().asSchema()); + assertFalse(schemaChanges.isEmpty()); + manager.applySchemaChange(table, schemaChanges); + } + + @Test + public void ignoresOlderRecordMissingColumn() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()), + Types.NestedField.optional(3, "email", Types.StringType.get()))); + + Schema olderSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + Record record = GenericRecord.create(olderSchema); + List schemaChanges = manager.checkSchemaChanges(table, record.struct().asSchema()); + assertTrue(schemaChanges.isEmpty()); + assertNotNull(catalog.loadTable(tableId).schema().findField("email")); + } + + @Test + public void ignoresOlderRecordWhenFieldAlreadyOptional() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.optional(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()))); + + Schema olderSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + Record record = GenericRecord.create(olderSchema); + List schemaChanges = manager.checkSchemaChanges(table, record.struct().asSchema()); + assertTrue(schemaChanges.isEmpty()); + assertTrue(catalog.loadTable(tableId).schema().findField("id").isOptional()); + } + + @Test + public void ignoresOlderRecordWhenTypeAlreadyPromoted() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get()))); + + Schema olderSchema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "name", Types.StringType.get())); + + Record record = GenericRecord.create(olderSchema); + List schemaChanges = manager.checkSchemaChanges(table, record.struct().asSchema()); + assertTrue(schemaChanges.isEmpty()); + assertEquals(Types.LongType.get(), catalog.loadTable(tableId).schema().findField("id").type()); + } + + @Test + public void doesNothingWhenSchemasMatch() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Schema schema = new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get())); + Table table = manager.getTableOrCreate(schema); + + Record record = GenericRecord.create(schema); + List schemaChanges = manager.checkSchemaChanges(table, record.struct().asSchema()); + assertTrue(schemaChanges.isEmpty()); + } + + @Test + public void skipsDuplicateNestedAdditions() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Table table = manager.getTableOrCreate(new Schema( + Types.NestedField.required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "address", Types.StructType.of( + Types.NestedField.optional(3, "street", Types.StringType.get()))))); + + List schemaChanges = List.of( + new IcebergTableManager.SchemaChange(IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, + "street", Types.StringType.get(), "address"), + new IcebergTableManager.SchemaChange(IcebergTableManager.SchemaChange.ChangeType.ADD_COLUMN, + "zipCode", Types.IntegerType.get(), "address")); + + manager.applySchemaChange(table, schemaChanges); + + Table updatedTable = catalog.loadTable(tableId); + Types.NestedField addressField = updatedTable.schema().findField("address"); + assertNotNull(addressField); + List nestedFields = addressField.type().asStructType().fields(); + assertEquals(2, nestedFields.size()); + assertNotNull(updatedTable.schema().findField("address.street")); + assertNotNull(updatedTable.schema().findField("address.zipCode")); + } + + @Test + public void skipsMakeOptionalAndPromoteWhenAlreadyApplied() { + TableIdentifier tableId = randomTableId(); + IcebergTableManager manager = newManager(tableId); + + Schema tableSchema = new Schema( + Types.NestedField.optional(1, "name", Types.StringType.get()), + Types.NestedField.required(2, "id", Types.LongType.get())); + + Table mockTable = mock(Table.class); + UpdateSchema mockUpdateSchema = mock(UpdateSchema.class); + when(mockTable.schema()).thenReturn(tableSchema); + when(mockTable.updateSchema()).thenReturn(mockUpdateSchema); + + List schemaChanges = List.of( + new IcebergTableManager.SchemaChange(IcebergTableManager.SchemaChange.ChangeType.MAKE_OPTIONAL, + "name", null, null), + new IcebergTableManager.SchemaChange(IcebergTableManager.SchemaChange.ChangeType.PROMOTE_TYPE, + "id", Types.LongType.get(), null)); + + manager.applySchemaChange(mockTable, schemaChanges); + + verify(mockUpdateSchema, never()).makeColumnOptional("name"); + verify(mockUpdateSchema, never()).updateColumn(eq("id"), any()); + verify(mockUpdateSchema).commit(); + } + + private Table applyChanges(IcebergTableManager manager, Table table, Schema newSchema) { + Record record = GenericRecord.create(newSchema); + List schemaChanges = manager.checkSchemaChanges(table, record.struct().asSchema()); + assertFalse(schemaChanges.isEmpty(), "Expected schema changes to be detected"); + manager.applySchemaChange(table, schemaChanges); + return catalog.loadTable(manager.tableId()); + } + + private IcebergTableManager newManager(TableIdentifier tableId) { + return new IcebergTableManager(catalog, tableId, mock(WorkerConfig.class)); + } + + private TableIdentifier randomTableId() { + return TableIdentifier.of("default", "table_" + ThreadLocalRandom.current().nextInt(1000, 10000)); + } + + private InMemoryCatalog initializeCatalog() { + InMemoryCatalog inMemoryCatalog = new InMemoryCatalog(); + inMemoryCatalog.initialize("test", ImmutableMap.of()); + return inMemoryCatalog; + } +} diff --git a/core/src/test/java/kafka/automq/table/worker/IcebergWriterSchemaEvolutionTest.java b/core/src/test/java/kafka/automq/table/worker/IcebergWriterSchemaEvolutionTest.java new file mode 100644 index 0000000000..5361032c27 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/IcebergWriterSchemaEvolutionTest.java @@ -0,0 +1,1454 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.process.DefaultRecordProcessor; +import kafka.automq.table.process.RecordProcessor; +import kafka.automq.table.process.convert.AvroRegistryConverter; +import kafka.automq.table.process.convert.StringConverter; + +import org.apache.kafka.common.header.Header; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.TimestampType; + +import com.google.common.collect.ImmutableMap; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.iceberg.Table; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.ThreadLocalRandom; + +import io.confluent.kafka.serializers.KafkaAvroDeserializer; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +class IcebergWriterSchemaEvolutionTest { + private static final String TOPIC = "test-topic"; + private static final byte MAGIC_BYTE = 0x0; + + @Mock + private KafkaAvroDeserializer kafkaAvroDeserializer; + + private InMemoryCatalog catalog; + private IcebergWriter writer; + private TableIdentifier tableId; + + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + catalog = new InMemoryCatalog(); + catalog.initialize("test", ImmutableMap.of()); + catalog.createNamespace(Namespace.of("default")); + String tableName = generateRandomTableName(); + tableId = TableIdentifier.of("default", tableName); + WorkerConfig config = mock(WorkerConfig.class); + when(config.partitionBy()).thenReturn(Collections.emptyList()); + IcebergTableManager tableManager = new IcebergTableManager(catalog, tableId, config); + + AvroRegistryConverter registryConverter = new AvroRegistryConverter(kafkaAvroDeserializer, null); + RecordProcessor processor = new DefaultRecordProcessor(TOPIC, StringConverter.INSTANCE, registryConverter); + writer = new IcebergWriter(tableManager, processor, config); + writer.setOffset(0, 0); + } + + @Test + void testAddRequiredFieldsInNestedStruct() throws IOException { + // v1: {id, user{name}} + Schema userV1 = Schema.createRecord("User", null, null, false); + userV1.setFields(Collections.singletonList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("name", "alice"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: Add required primitive field in nested struct (age) + // Add required nested struct field (profile) + // Add required list field in nested struct (hobbies) + Schema profileV2 = Schema.createRecord("Profile", null, null, false); + profileV2.setFields(Arrays.asList( + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("level", Schema.create(Schema.Type.INT), null, null))); + + Schema userV2 = Schema.createRecord("User", null, null, false); + Schema hobbiesSchema = Schema.createArray(Schema.create(Schema.Type.STRING)); + userV2.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("age", Schema.create(Schema.Type.INT), null, null), + new Schema.Field("profile", profileV2, null, null), + new Schema.Field("hobbies", hobbiesSchema, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord profileRecordV2 = new GenericData.Record(profileV2); + profileRecordV2.put("city", "Shanghai"); + profileRecordV2.put("level", 5); + + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("name", "bob"); + userRecordV2.put("age", 30); + userRecordV2.put("profile", profileRecordV2); + userRecordV2.put("hobbies", Arrays.asList("reading", "coding")); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify nested primitive field + assertNotNull(table.schema().findField("_kafka_value.user.age")); + // Verify nested struct field + assertNotNull(table.schema().findField("_kafka_value.user.profile")); + assertNotNull(table.schema().findField("_kafka_value.user.profile.city")); + // Verify nested list field + assertNotNull(table.schema().findField("_kafka_value.user.hobbies")); + } + + @Test + void testAddRequiredFieldsInCollectionElement() throws IOException { + // v1: {id, addresses: list<{street}>} + Schema addressV1 = Schema.createRecord("Address", null, null, false); + addressV1.setFields(Collections.singletonList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(addressV1); + Schema optionalListV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV1, null, null))); + + GenericRecord addressRecordV1 = new GenericData.Record(addressV1); + addressRecordV1.put("street", "Main St"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("addresses", Collections.singletonList(addressRecordV1)); + + // v2: Add required primitive field in list element (zipCode, floor) + // Add required nested struct field in list element (location) + Schema locationV2 = Schema.createRecord("Location", null, null, false); + locationV2.setFields(Arrays.asList( + new Schema.Field("lat", Schema.create(Schema.Type.DOUBLE), null, null), + new Schema.Field("lng", Schema.create(Schema.Type.DOUBLE), null, null))); + + Schema addressV2 = Schema.createRecord("Address", null, null, false); + addressV2.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("zipCode", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("floor", Schema.create(Schema.Type.INT), null, null), + new Schema.Field("location", locationV2, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV2, null, null))); + + GenericRecord locationRecordV2 = new GenericData.Record(locationV2); + locationRecordV2.put("lat", 39.9042); + locationRecordV2.put("lng", 116.4074); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("street", "Second St"); + addressRecordV2.put("zipCode", "100000"); + addressRecordV2.put("floor", 5); + addressRecordV2.put("location", locationRecordV2); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify primitive fields in list element + assertNotNull(table.schema().findField("_kafka_value.addresses.element.zipCode")); + assertNotNull(table.schema().findField("_kafka_value.addresses.element.floor")); + // Verify struct field in list element + assertNotNull(table.schema().findField("_kafka_value.addresses.element.location")); + assertNotNull(table.schema().findField("_kafka_value.addresses.element.location.lat")); + } + + @Test + void testAddOptionalCollection() throws IOException { + // v1: {id, name} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("name", "alice"); + + // v2: {id, name, tags: list} + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listSchema = Schema.createArray(Schema.create(Schema.Type.STRING)); + Schema optionalList = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listSchema)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("tags", optionalList, null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("name", "bob"); + avroRecordV2.put("tags", Arrays.asList("tag1", "tag2")); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.tags")); + assertEquals(true, table.schema().findField("_kafka_value.tags").type().isListType()); + } + + + @Test + void testAddRequiredCollections() throws IOException { + // v1: {id, name} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("name", "alice"); + + // v2: Add list (tags) + // Add list (addresses) + // Add map (scores) + // Add map (locations) + Schema addressSchema = Schema.createRecord("Address", null, null, false); + addressSchema.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null))); + + Schema locationSchema = Schema.createRecord("Location", null, null, false); + locationSchema.setFields(Arrays.asList( + new Schema.Field("lat", Schema.create(Schema.Type.DOUBLE), null, null), + new Schema.Field("lng", Schema.create(Schema.Type.DOUBLE), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema tagsList = Schema.createArray(Schema.create(Schema.Type.STRING)); + Schema addressesList = Schema.createArray(addressSchema); + Schema scoresMap = Schema.createMap(Schema.create(Schema.Type.INT)); + Schema locationsMap = Schema.createMap(locationSchema); + + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("tags", tagsList, null, null), + new Schema.Field("addresses", addressesList, null, null), + new Schema.Field("scores", scoresMap, null, null), + new Schema.Field("locations", locationsMap, null, null))); + + GenericRecord addressRecord = new GenericData.Record(addressSchema); + addressRecord.put("street", "Main St"); + addressRecord.put("city", "Beijing"); + + GenericRecord locationRecord = new GenericData.Record(locationSchema); + locationRecord.put("lat", 39.9042); + locationRecord.put("lng", 116.4074); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("name", "bob"); + avroRecordV2.put("tags", Arrays.asList("tag1", "tag2")); + avroRecordV2.put("addresses", Collections.singletonList(addressRecord)); + java.util.Map scoresMapData = new java.util.HashMap<>(); + scoresMapData.put("math", 95); + avroRecordV2.put("scores", scoresMapData); + java.util.Map locationsMapData = new java.util.HashMap<>(); + locationsMapData.put("home", locationRecord); + avroRecordV2.put("locations", locationsMapData); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify list + assertNotNull(table.schema().findField("_kafka_value.tags")); + assertEquals(true, table.schema().findField("_kafka_value.tags").type().isListType()); + // Verify list + assertNotNull(table.schema().findField("_kafka_value.addresses")); + assertNotNull(table.schema().findField("_kafka_value.addresses.element.street")); + // Verify map + assertNotNull(table.schema().findField("_kafka_value.scores")); + assertEquals(true, table.schema().findField("_kafka_value.scores").type().isMapType()); + // Verify map + assertNotNull(table.schema().findField("_kafka_value.locations")); + assertNotNull(table.schema().findField("_kafka_value.locations.value.lat")); + } + + + @Test + void testAddRequiredFieldInCollectionElement() throws IOException { + // v1: addresses: list<{street}> + Schema addressV1 = Schema.createRecord("Address", null, null, false); + addressV1.setFields(Collections.singletonList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(addressV1); + Schema optionalListV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV1, null, null))); + + GenericRecord addressRecordV1 = new GenericData.Record(addressV1); + addressRecordV1.put("street", "Main St"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("addresses", Collections.singletonList(addressRecordV1)); + + // v2: addresses: list<{street, required city}> + Schema addressV2 = Schema.createRecord("Address", null, null, false); + addressV2.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV2, null, null))); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("street", "Second St"); + addressRecordV2.put("city", "Beijing"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.addresses.element.city")); + assertEquals(Types.StringType.get(), table.schema().findField("_kafka_value.addresses.element.city").type()); + } + + @Test + void testAddRequiredFieldInMapValueStruct() throws IOException { + // v1: locations: map + Schema locationV1 = Schema.createRecord("Location", null, null, false); + locationV1.setFields(Collections.singletonList( + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema mapV1 = Schema.createMap(locationV1); + Schema optionalMapV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("locations", optionalMapV1, null, null))); + + GenericRecord locationRecordV1 = new GenericData.Record(locationV1); + locationRecordV1.put("city", "Shanghai"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + java.util.Map mapData1 = new java.util.HashMap<>(); + mapData1.put("home", locationRecordV1); + avroRecordV1.put("locations", mapData1); + + // v2: locations: map + Schema locationV2 = Schema.createRecord("Location", null, null, false); + locationV2.setFields(Arrays.asList( + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("country", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema mapV2 = Schema.createMap(locationV2); + Schema optionalMapV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("locations", optionalMapV2, null, null))); + + GenericRecord locationRecordV2 = new GenericData.Record(locationV2); + locationRecordV2.put("city", "Beijing"); + locationRecordV2.put("country", "China"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + java.util.Map mapData2 = new java.util.HashMap<>(); + mapData2.put("work", locationRecordV2); + avroRecordV2.put("locations", mapData2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.locations.value.country")); + assertEquals(Types.StringType.get(), table.schema().findField("_kafka_value.locations.value.country").type()); + } + + + // ========== Add Optional Field Tests ========== + + @Test + void testAddOptionalFieldInNestedStruct() throws IOException { + // v1: user{name} + Schema userV1 = Schema.createRecord("User", null, null, false); + userV1.setFields(Collections.singletonList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("name", "alice"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: user{name, email} + Schema userV2 = Schema.createRecord("User", null, null, false); + Schema optionalEmail = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + userV2.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("email", optionalEmail, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("name", "bob"); + userRecordV2.put("email", "bob@example.com"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.user.email")); + } + + + @Test + void testDropRequiredCollection() throws IOException { + // v1: {id, tags: list} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listSchema = Schema.createArray(Schema.create(Schema.Type.STRING)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("tags", listSchema, null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("tags", Arrays.asList("tag1", "tag2")); + + // v2: {id} - dropped tags + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Collections.singletonList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.tags")); + assertEquals(false, table.schema().findField("_kafka_value.tags").isRequired()); + } + + @Test + void testDropRequiredFieldInNestedStruct() throws IOException { + // v1: user{name, email, age} + Schema userV1 = Schema.createRecord("User", null, null, false); + userV1.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("email", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("age", Schema.create(Schema.Type.INT), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("name", "alice"); + userRecordV1.put("email", "alice@example.com"); + userRecordV1.put("age", 25); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: user{name, email} - dropped age + Schema userV2 = Schema.createRecord("User", null, null, false); + userV2.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("email", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("name", "bob"); + userRecordV2.put("email", "bob@example.com"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.user.age")); + assertEquals(false, table.schema().findField("_kafka_value.user.age").isRequired()); + } + + @Test + void testMakeRequiredFieldOptionalInNestedStruct() throws IOException { + // v1: user{required name, required age} + Schema userV1 = Schema.createRecord("User", null, null, false); + userV1.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("age", Schema.create(Schema.Type.INT), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("name", "alice"); + userRecordV1.put("age", 25); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: user{required name, optional age} + Schema userV2 = Schema.createRecord("User", null, null, false); + Schema optionalAge = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.INT))); + userV2.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("age", optionalAge, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("name", "bob"); + userRecordV2.put("age", null); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertEquals(false, table.schema().findField("_kafka_value.user.age").isRequired()); + } + + @Test + void testPromoteFieldTypeInNestedStruct() throws IOException { + // v1: user{age: int} + Schema userV1 = Schema.createRecord("User", null, null, false); + userV1.setFields(Collections.singletonList( + new Schema.Field("age", Schema.create(Schema.Type.INT), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("age", 25); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: user{age: long} + Schema userV2 = Schema.createRecord("User", null, null, false); + userV2.setFields(Collections.singletonList( + new Schema.Field("age", Schema.create(Schema.Type.LONG), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("age", 30L); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertEquals(Types.LongType.get(), table.schema().findField("_kafka_value.user.age").type()); + } + + // ========== Collection Field Tests ========== + + + + @Test + void testAddOptionalMapCollection() throws IOException { + // v1: {id, name} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("name", "alice"); + + // v2: {id, name, metadata: map} + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema mapSchema = Schema.createMap(Schema.create(Schema.Type.STRING)); + Schema optionalMap = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapSchema)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("metadata", optionalMap, null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("name", "bob"); + java.util.Map metadataMap = new java.util.HashMap<>(); + metadataMap.put("key1", "value1"); + metadataMap.put("key2", "value2"); + avroRecordV2.put("metadata", metadataMap); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.metadata")); + assertEquals(true, table.schema().findField("_kafka_value.metadata").type().isMapType()); + } + + + @Test + void testAddOptionalFieldInCollectionElement() throws IOException { + // v1: addresses: list<{street}> + Schema addressV1 = Schema.createRecord("Address", null, null, false); + addressV1.setFields(Collections.singletonList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(addressV1); + Schema optionalListV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV1, null, null))); + + GenericRecord addressRecordV1 = new GenericData.Record(addressV1); + addressRecordV1.put("street", "Main St"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("addresses", Collections.singletonList(addressRecordV1)); + + // v2: addresses: list<{street, zipCode}> + Schema addressV2 = Schema.createRecord("Address", null, null, false); + Schema optionalZip = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + addressV2.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("zipCode", optionalZip, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV2, null, null))); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("street", "Second St"); + addressRecordV2.put("zipCode", "12345"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.addresses.element.zipCode")); + } + + @Test + void testDropRequiredFieldInMapValueStruct() throws IOException { + // v1: attributes: map + Schema locationV1 = Schema.createRecord("Location", null, null, false); + locationV1.setFields(Arrays.asList( + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("country", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema mapV1 = Schema.createMap(locationV1); + Schema optionalMapV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("attributes", optionalMapV1, null, null))); + + GenericRecord locationRecordV1 = new GenericData.Record(locationV1); + locationRecordV1.put("city", "Beijing"); + locationRecordV1.put("country", "China"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + java.util.Map mapData1 = new java.util.HashMap<>(); + mapData1.put("home", locationRecordV1); + avroRecordV1.put("attributes", mapData1); + + // v2: attributes: map - dropped country + Schema locationV2 = Schema.createRecord("Location", null, null, false); + locationV2.setFields(Collections.singletonList( + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema mapV2 = Schema.createMap(locationV2); + Schema optionalMapV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("attributes", optionalMapV2, null, null))); + + GenericRecord locationRecordV2 = new GenericData.Record(locationV2); + locationRecordV2.put("city", "Shanghai"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + java.util.Map mapData2 = new java.util.HashMap<>(); + mapData2.put("work", locationRecordV2); + avroRecordV2.put("attributes", mapData2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.attributes.value.country")); + assertEquals(false, table.schema().findField("_kafka_value.attributes.value.country").isRequired()); + } + + @Test + void testPromoteFieldTypeInCollectionElement() throws IOException { + // v1: addresses: list<{zip: int}> + Schema addressV1 = Schema.createRecord("Address", null, null, false); + addressV1.setFields(Collections.singletonList( + new Schema.Field("zip", Schema.create(Schema.Type.INT), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(addressV1); + Schema optionalListV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV1, null, null))); + + GenericRecord addressRecordV1 = new GenericData.Record(addressV1); + addressRecordV1.put("zip", 12345); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("addresses", Collections.singletonList(addressRecordV1)); + + // v2: addresses: list<{zip: long}> + Schema addressV2 = Schema.createRecord("Address", null, null, false); + addressV2.setFields(Collections.singletonList( + new Schema.Field("zip", Schema.create(Schema.Type.LONG), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV2, null, null))); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("zip", 67890L); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertEquals(Types.LongType.get(), table.schema().findField("_kafka_value.addresses.element.zip").type()); + } + + // ========== Drop Required Field Tests ========== + + @Test + void testDropRequiredFieldInCollectionElement() throws IOException { + // v1: addresses: list<{street, city, zipCode}> + Schema addressV1 = Schema.createRecord("Address", null, null, false); + addressV1.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("zipCode", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(addressV1); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", listV1, null, null))); + + GenericRecord addressRecordV1 = new GenericData.Record(addressV1); + addressRecordV1.put("street", "Main St"); + addressRecordV1.put("city", "Beijing"); + addressRecordV1.put("zipCode", "100000"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("addresses", Collections.singletonList(addressRecordV1)); + + // v2: addresses: list<{street, city}> - dropped required zipCode + Schema addressV2 = Schema.createRecord("Address", null, null, false); + addressV2.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", listV2, null, null))); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("street", "Second St"); + addressRecordV2.put("city", "Shanghai"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify that the dropped required field is now optional + assertNotNull(table.schema().findField("_kafka_value.addresses.element.zipCode")); + assertEquals(false, table.schema().findField("_kafka_value.addresses.element.zipCode").isRequired()); + } + + // ========== Drop Optional Field Tests ========== + + @Test + void testDropOptionalFieldInNestedStruct() throws IOException { + // v1: user{name, email, phone} + Schema userV1 = Schema.createRecord("User", null, null, false); + Schema optionalEmail = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + Schema optionalPhone = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + userV1.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("email", optionalEmail, null, null), + new Schema.Field("phone", optionalPhone, null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("name", "alice"); + userRecordV1.put("email", "alice@example.com"); + userRecordV1.put("phone", "123-456-7890"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: user{name, email} - dropped optional phone + Schema userV2 = Schema.createRecord("User", null, null, false); + userV2.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("email", optionalEmail, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("name", "bob"); + userRecordV2.put("email", "bob@example.com"); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify that the dropped optional field still exists and remains optional + assertNotNull(table.schema().findField("_kafka_value.user.phone")); + assertEquals(false, table.schema().findField("_kafka_value.user.phone").isRequired()); + } + + @Test + void testDropOptionalCollections() throws IOException { + // v1: {id, tags: optional list, metadata: optional map} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listSchema = Schema.createArray(Schema.create(Schema.Type.STRING)); + Schema optionalList = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listSchema)); + Schema mapSchema = Schema.createMap(Schema.create(Schema.Type.STRING)); + Schema optionalMap = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapSchema)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("tags", optionalList, null, null), + new Schema.Field("metadata", optionalMap, null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("tags", Arrays.asList("tag1", "tag2")); + java.util.Map metadataMap = new java.util.HashMap<>(); + metadataMap.put("key1", "value1"); + avroRecordV1.put("metadata", metadataMap); + + // v2: {id} - dropped both optional collections + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Collections.singletonList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify that dropped optional collections still exist and remain optional + assertNotNull(table.schema().findField("_kafka_value.tags")); + assertEquals(false, table.schema().findField("_kafka_value.tags").isRequired()); + assertNotNull(table.schema().findField("_kafka_value.metadata")); + assertEquals(false, table.schema().findField("_kafka_value.metadata").isRequired()); + } + + // ========== Make Field Optional Tests ========== + + @Test + void testMakeRequiredFieldOptional() throws IOException { + // v1: {id, required email} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("email", Schema.create(Schema.Type.STRING), null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("email", "alice@example.com"); + + // v2: {id, optional email} + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema optionalEmail = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("email", optionalEmail, null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("email", null); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.email")); + assertEquals(false, table.schema().findField("_kafka_value.email").isRequired()); + } + + @Test + void testMakeRequiredFieldOptionalInCollectionElement() throws IOException { + // v1: addresses: list<{street, required city}> + Schema addressV1 = Schema.createRecord("Address", null, null, false); + addressV1.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("city", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(addressV1); + Schema optionalListV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV1, null, null))); + + GenericRecord addressRecordV1 = new GenericData.Record(addressV1); + addressRecordV1.put("street", "Main St"); + addressRecordV1.put("city", "Beijing"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("addresses", Collections.singletonList(addressRecordV1)); + + // v2: addresses: list<{street, optional city}> + Schema addressV2 = Schema.createRecord("Address", null, null, false); + Schema optionalCity = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + addressV2.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("city", optionalCity, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("addresses", optionalListV2, null, null))); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("street", "Second St"); + addressRecordV2.put("city", null); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify that the field in collection element is now optional + assertNotNull(table.schema().findField("_kafka_value.addresses.element.city")); + assertEquals(false, table.schema().findField("_kafka_value.addresses.element.city").isRequired()); + } + + @Test + void testMakeRequiredCollectionOptional() throws IOException { + // v1: {id, required tags: list} + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(Schema.create(Schema.Type.STRING)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("tags", listV1, null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("tags", Arrays.asList("tag1", "tag2")); + + // v2: {id, optional tags: list} + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(Schema.create(Schema.Type.STRING)); + Schema optionalList = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("tags", optionalList, null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("tags", null); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.tags")); + assertEquals(false, table.schema().findField("_kafka_value.tags").isRequired()); + } + + // ========== Promote Field Type Tests ========== + + @Test + void testPromoteCollectionElementType() throws IOException { + // v1: scores: list + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema listV1 = Schema.createArray(Schema.create(Schema.Type.INT)); + Schema optionalListV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("scores", optionalListV1, null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("scores", Arrays.asList(90, 85, 95)); + + // v2: scores: list + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema listV2 = Schema.createArray(Schema.create(Schema.Type.LONG)); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("scores", optionalListV2, null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("scores", Arrays.asList(100L, 95L, 98L)); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify that the collection element type cannot be promoted from int to long + assertNotNull(table.schema().findField("_kafka_value.scores")); + assertEquals(true, table.schema().findField("_kafka_value.scores").type().isListType()); + assertEquals(Types.IntegerType.get(), + table.schema().findField("_kafka_value.scores").type().asListType().elementType()); + } + + @Test + void testPromoteMapValueType() throws IOException { + // v1: metadata: map + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + Schema mapV1 = Schema.createMap(Schema.create(Schema.Type.INT)); + Schema optionalMapV1 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapV1)); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("metadata", optionalMapV1, null, null))); + + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + java.util.Map mapData1 = new java.util.HashMap<>(); + mapData1.put("score", 90); + mapData1.put("rank", 5); + avroRecordV1.put("metadata", mapData1); + + // v2: metadata: map + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + Schema mapV2 = Schema.createMap(Schema.create(Schema.Type.LONG)); + Schema optionalMapV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), mapV2)); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("metadata", optionalMapV2, null, null))); + + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + java.util.Map mapData2 = new java.util.HashMap<>(); + mapData2.put("score", 100L); + mapData2.put("rank", 1L); + avroRecordV2.put("metadata", mapData2); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + + Table table = catalog.loadTable(tableId); + // Verify that the map value type cannot be promoted from int to long + assertNotNull(table.schema().findField("_kafka_value.metadata")); + assertEquals(true, table.schema().findField("_kafka_value.metadata").type().isMapType()); + assertEquals(Types.IntegerType.get(), + table.schema().findField("_kafka_value.metadata").type().asMapType().valueType()); + } + + // ========== Complex Mixed Scenarios ========== + + @Test + void testComplexNestedAndCollectionEvolution() throws IOException { + // v1: {id, user{name}} + Schema userV1 = Schema.createRecord("User", null, null, false); + userV1.setFields(Collections.singletonList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null))); + + Schema avroSchemaV1 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV1.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV1, null, null))); + + GenericRecord userRecordV1 = new GenericData.Record(userV1); + userRecordV1.put("name", "alice"); + GenericRecord avroRecordV1 = new GenericData.Record(avroSchemaV1); + avroRecordV1.put("id", 1L); + avroRecordV1.put("user", userRecordV1); + + // v2: {id, user{name, addresses: list<{street}>}} + Schema addressV2 = Schema.createRecord("Address", null, null, false); + addressV2.setFields(Collections.singletonList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null))); + + Schema userV2 = Schema.createRecord("User", null, null, false); + Schema listV2 = Schema.createArray(addressV2); + Schema optionalListV2 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV2)); + userV2.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("addresses", optionalListV2, null, null))); + + Schema avroSchemaV2 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV2.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV2, null, null))); + + GenericRecord addressRecordV2 = new GenericData.Record(addressV2); + addressRecordV2.put("street", "Main St"); + GenericRecord userRecordV2 = new GenericData.Record(userV2); + userRecordV2.put("name", "bob"); + userRecordV2.put("addresses", Collections.singletonList(addressRecordV2)); + GenericRecord avroRecordV2 = new GenericData.Record(avroSchemaV2); + avroRecordV2.put("id", 2L); + avroRecordV2.put("user", userRecordV2); + + // v3: {id, user{name, addresses: list<{street, zipCode}>}} + Schema addressV3 = Schema.createRecord("Address", null, null, false); + Schema optionalZip = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), Schema.create(Schema.Type.STRING))); + addressV3.setFields(Arrays.asList( + new Schema.Field("street", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("zipCode", optionalZip, null, null))); + + Schema userV3 = Schema.createRecord("User", null, null, false); + Schema listV3 = Schema.createArray(addressV3); + Schema optionalListV3 = Schema.createUnion(Arrays.asList(Schema.create(Schema.Type.NULL), listV3)); + userV3.setFields(Arrays.asList( + new Schema.Field("name", Schema.create(Schema.Type.STRING), null, null), + new Schema.Field("addresses", optionalListV3, null, null))); + + Schema avroSchemaV3 = Schema.createRecord("TestRecord", null, null, false); + avroSchemaV3.setFields(Arrays.asList( + new Schema.Field("id", Schema.create(Schema.Type.LONG), null, null), + new Schema.Field("user", userV3, null, null))); + + GenericRecord addressRecordV3 = new GenericData.Record(addressV3); + addressRecordV3.put("street", "Second St"); + addressRecordV3.put("zipCode", "12345"); + GenericRecord userRecordV3 = new GenericData.Record(userV3); + userRecordV3.put("name", "charlie"); + userRecordV3.put("addresses", Collections.singletonList(addressRecordV3)); + GenericRecord avroRecordV3 = new GenericData.Record(avroSchemaV3); + avroRecordV3.put("id", 3L); + avroRecordV3.put("user", userRecordV3); + + when(kafkaAvroDeserializer.deserialize(anyString(), any(), any(ByteBuffer.class))) + .thenReturn(avroRecordV1) + .thenReturn(avroRecordV2) + .thenReturn(avroRecordV3); + + writer.write(0, createMockKafkaRecord(1, 0)); + writer.write(0, createMockKafkaRecord(2, 1)); + writer.write(0, createMockKafkaRecord(3, 2)); + + Table table = catalog.loadTable(tableId); + assertNotNull(table.schema().findField("_kafka_value.user.addresses")); + assertNotNull(table.schema().findField("_kafka_value.user.addresses.element.street")); + assertNotNull(table.schema().findField("_kafka_value.user.addresses.element.zipCode")); + } + + // ========== Helper Methods ========== + + private String generateRandomTableName() { + return "test_table_" + ThreadLocalRandom.current().nextInt(100000, 999999); + } + + private Record createMockKafkaRecord(int schemaId, int offset) { + ByteBuffer value = ByteBuffer.allocate(5 + 10); // 1 byte magic + 4 bytes schema ID + some space for data + value.put(MAGIC_BYTE); + value.putInt(schemaId); + // Add some dummy data + value.put("test".getBytes()); + + return new Record() { + @Override + public long offset() { + return offset; + } + + @Override + public int sequence() { + return 0; + } + + @Override + public int sizeInBytes() { + return value.array().length; + } + + @Override + public long timestamp() { + return System.currentTimeMillis(); + } + + @Override + public void ensureValid() { + } + + @Override + public int keySize() { + return 0; + } + + @Override + public boolean hasKey() { + return false; + } + + @Override + public ByteBuffer key() { + return null; + } + + @Override + public int valueSize() { + return value.array().length; + } + + @Override + public boolean hasValue() { + return true; + } + + @Override + public ByteBuffer value() { + return ByteBuffer.wrap(value.array()); + } + + @Override + public boolean hasMagic(byte b) { + return b == MAGIC_BYTE; + } + + @Override + public boolean isCompressed() { + return false; + } + + @Override + public boolean hasTimestampType(TimestampType timestampType) { + return false; + } + + @Override + public Header[] headers() { + return new Header[0]; + } + }; + } +} diff --git a/core/src/test/java/kafka/automq/table/worker/MemoryWriter.java b/core/src/test/java/kafka/automq/table/worker/MemoryWriter.java new file mode 100644 index 0000000000..7d10810bf0 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/MemoryWriter.java @@ -0,0 +1,61 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.process.DefaultRecordProcessor; +import kafka.automq.table.process.convert.RawConverter; +import kafka.automq.table.process.exception.RecordProcessorException; + +import org.apache.kafka.common.record.Record; + +import com.google.common.collect.ImmutableMap; + +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.Namespace; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.inmemory.InMemoryCatalog; +import org.junit.jupiter.api.Tag; + +import java.io.IOException; +import java.util.LinkedList; +import java.util.List; + +@Tag("S3Unit") +public class MemoryWriter extends IcebergWriter { + List records = new LinkedList<>(); + + public MemoryWriter(WorkerConfig config) { + super(new IcebergTableManager(catalog(), TableIdentifier.parse("default.test"), config), new DefaultRecordProcessor("test", new RawConverter(), new RawConverter()), config); + } + + @Override + protected boolean write0(int partition, Record kafkaRecord) throws IOException, RecordProcessorException { + super.write0(partition, kafkaRecord); + records.add(kafkaRecord); + return true; + } + + private static Catalog catalog() { + InMemoryCatalog catalog = new InMemoryCatalog(); + catalog.initialize("test", ImmutableMap.of()); + catalog.createNamespace(Namespace.of("default")); + return catalog; + } +} diff --git a/core/src/test/java/kafka/automq/table/worker/RecordIteratorTest.java b/core/src/test/java/kafka/automq/table/worker/RecordIteratorTest.java new file mode 100644 index 0000000000..3dcfe4ae92 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/RecordIteratorTest.java @@ -0,0 +1,143 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.record.ControlRecordType; +import org.apache.kafka.common.record.EndTransactionMarker; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.MemoryRecordsBuilder; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.RecordBatch; +import org.apache.kafka.common.record.Records; +import org.apache.kafka.common.record.SimpleRecord; +import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.common.utils.BufferSupplier; +import org.apache.kafka.storage.internals.log.FetchDataInfo; + +import com.automq.stream.utils.MockTime; +import com.automq.stream.utils.Time; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +@Tag("S3Unit") +public class RecordIteratorTest { + private static final String TOPIC_NAME = "test"; + + private static final TopicPartition TP = new TopicPartition(TOPIC_NAME, 0); + private static final long PRODUCER_ID = 1000L; + private static final short PRODUCER_EPOCH = 0; + + @Test + public void testAbortedTransactionRecordsRemoved() { + Records abortRawRecords = newTranscactionalRecords(ControlRecordType.ABORT, 0, 10); + Records commitrawRecords = newTranscactionalRecords(ControlRecordType.COMMIT, 11, 1); + ByteBuf buf = Unpooled.buffer(); + buf.writeBytes(((MemoryRecords) abortRawRecords).buffer()); + buf.writeBytes(((MemoryRecords) commitrawRecords).buffer()); + Records rawRecords = MemoryRecords.readableRecords(buf.nioBuffer()); + + RecordsIterator it = new RecordsIterator(0, new FetchDataInfo(null, rawRecords, false, Optional.of(newAbortedTransactions())), BufferSupplier.NO_CACHING); + Assertions.assertTrue(it.hasNext()); + Record record = it.next(); + Assertions.assertEquals("value11", new String(buf2bytes(record.value()))); + Assertions.assertFalse(it.hasNext()); + Assertions.assertEquals(13, it.nextOffset()); + } + + @Test + public void testCommittedTransactionRecordsIncluded() { + int numRecords = 10; + Records rawRecords = newTranscactionalRecords(ControlRecordType.COMMIT, 0, numRecords); + + RecordsIterator it = new RecordsIterator(0, new FetchDataInfo(null, rawRecords, false, Optional.empty()), BufferSupplier.NO_CACHING); + for (int i = 0; i < 10; i++) { + Assertions.assertTrue(it.hasNext()); + it.next(); + } + Assertions.assertFalse(it.hasNext()); + Assertions.assertEquals(11, it.nextOffset()); + } + + private Records newTranscactionalRecords(ControlRecordType controlRecordType, long baseOffset, int numRecords) { + Time time = new MockTime(); + ByteBuffer buffer = ByteBuffer.allocate(1024); + + try (MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, + RecordBatch.CURRENT_MAGIC_VALUE, + Compression.NONE, + TimestampType.CREATE_TIME, + baseOffset, + time.milliseconds(), + PRODUCER_ID, + PRODUCER_EPOCH, + 0, + true, + RecordBatch.NO_PARTITION_LEADER_EPOCH)) { + for (int i = 0; i < numRecords; i++) + builder.append(new SimpleRecord(time.milliseconds(), "key".getBytes(), ("value" + (baseOffset + i)).getBytes())); + + builder.build(); + } + + writeTransactionMarker(buffer, controlRecordType, baseOffset + numRecords, time); + buffer.flip(); + + return MemoryRecords.readableRecords(buffer); + } + + private void writeTransactionMarker(ByteBuffer buffer, + ControlRecordType controlRecordType, + long offset, + Time time) { + MemoryRecords.writeEndTransactionalMarker(buffer, + offset, + time.milliseconds(), + 0, + PRODUCER_ID, + PRODUCER_EPOCH, + new EndTransactionMarker(controlRecordType, 0)); + } + + private List newAbortedTransactions() { + FetchResponseData.AbortedTransaction abortedTransaction = new FetchResponseData.AbortedTransaction(); + abortedTransaction.setFirstOffset(0); + abortedTransaction.setProducerId(PRODUCER_ID); + return Collections.singletonList(abortedTransaction); + } + + static byte[] buf2bytes(ByteBuffer buf) { + byte[] bytes = new byte[buf.remaining()]; + buf.duplicate().get(bytes); + return bytes; + } +} diff --git a/core/src/test/java/kafka/automq/table/worker/TestUtils.java b/core/src/test/java/kafka/automq/table/worker/TestUtils.java new file mode 100644 index 0000000000..e88fef1b8c --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/TestUtils.java @@ -0,0 +1,59 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.cluster.Partition; +import kafka.log.UnifiedLog; + +import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.SimpleRecord; + +import org.junit.jupiter.api.Tag; + +import scala.Option; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +public class TestUtils { + public static Partition partition(int partitionId, int epoch) { + Partition partition = mock(Partition.class); + when(partition.partitionId()).thenReturn(partitionId); + when(partition.getLeaderEpoch()).thenReturn(epoch); + UnifiedLog log = mock(UnifiedLog.class); + when(partition.log()).thenReturn(Option.apply(log)); + return partition; + } + + public static MemoryRecords newMemoryRecord(long offset, int count) { + return newMemoryRecord(offset, count, 1); + } + + public static MemoryRecords newMemoryRecord(long offset, int count, int recordSize) { + SimpleRecord[] records = new SimpleRecord[count]; + for (int i = 0; i < count; i++) { + records[i] = new SimpleRecord(new byte[recordSize]); + } + + return MemoryRecords.withRecords(offset, Compression.NONE, records); + } +} diff --git a/core/src/test/java/kafka/automq/table/worker/TopicPartitionsWorkerTest.java b/core/src/test/java/kafka/automq/table/worker/TopicPartitionsWorkerTest.java new file mode 100644 index 0000000000..28daf48853 --- /dev/null +++ b/core/src/test/java/kafka/automq/table/worker/TopicPartitionsWorkerTest.java @@ -0,0 +1,471 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.table.worker; + +import kafka.automq.table.Channel; +import kafka.automq.table.events.CommitRequest; +import kafka.automq.table.events.CommitResponse; +import kafka.automq.table.events.Errors; +import kafka.automq.table.events.Event; +import kafka.automq.table.events.WorkerOffset; +import kafka.automq.table.worker.TopicPartitionsWorker.OffsetBound; +import kafka.automq.table.worker.TopicPartitionsWorker.SyncTask; +import kafka.cluster.Partition; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.storage.internals.log.FetchDataInfo; + +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.threads.EventLoop; + +import org.apache.iceberg.PartitionSpec; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.parallel.Execution; +import org.mockito.ArgumentCaptor; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicReference; + +import static kafka.automq.table.worker.TestUtils.newMemoryRecord; +import static kafka.automq.table.worker.TestUtils.partition; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.parallel.ExecutionMode.SAME_THREAD; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +@Timeout(10) +@Execution(SAME_THREAD) +public class TopicPartitionsWorkerTest { + private static final String TOPIC = "TP_TEST"; + + TopicPartitionsWorker topicPartitionsWorker; + WriterFactory writerFactory; + List writers; + Channel channel; + WorkerConfig config; + + EventLoop eventLoop = new EventLoop("test"); + ExecutorService flushExecutor = Threads.newFixedThreadPoolWithMonitor(1, "test-flush", true, LoggerFactory.getLogger(TopicPartitionsWorkerTest.class)); + EventLoops eventLoops; + Semaphore commitLimiter; + + volatile double avgRecordRecord = 1; + + @BeforeEach + public void setup() { + config = mock(WorkerConfig.class); + when(config.microSyncBatchSize()).thenReturn(1024); + when(config.incrementSyncThreshold()).thenReturn(1024L); + + channel = mock(Channel.class); + when(channel.asyncSend(eq(TOPIC), any())).thenReturn(CompletableFuture.completedFuture(null)); + + writerFactory = mock(WriterFactory.class); + when(writerFactory.partitionSpec()).thenReturn(PartitionSpec.unpartitioned()); + writers = new ArrayList<>(); + when(writerFactory.newWriter()).thenAnswer(invocation -> { + MemoryWriter writer = new MemoryWriter(config); + writers.add(writer); + return writer; + }); + + avgRecordRecord = 1; + eventLoops = new EventLoops(new EventLoop[] {eventLoop, new EventLoop("worker-test")}); + commitLimiter = new Semaphore(4); + topicPartitionsWorker = new TopicPartitionsWorker(TOPIC, config, writerFactory, channel, eventLoop, eventLoops, flushExecutor, commitLimiter); + } + + @Test + public void testCommit() throws Exception { + Partition p1 = partition(1, 1); + + Partition p2 = partition(2, 3); + + // 1. commit with p0 missing, p1 mismatch, p2 match + { + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, newMemoryRecord(30, 3)); + when(p2.log().get().readAsync(eq(30L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().highWatermark()).thenReturn(33L); + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(0, 1, 10), + new WorkerOffset(1, 2, 20), + new WorkerOffset(2, 3, 30))); + AtomicReference> commitCf = new AtomicReference<>(); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p1); + topicPartitionsWorker.add(p2); + commitCf.set(topicPartitionsWorker.commit(commitRequest)); + }).get(); + commitCf.get().get(); + + verify(p2.log().get(), times(1)).readAsync(anyLong(), anyInt(), any(), anyBoolean()); + + ArgumentCaptor ac1 = ArgumentCaptor.forClass(Event.class); + verify(channel, times(2)).asyncSend(eq(TOPIC), ac1.capture()); + CommitResponse resp1 = ac1.getAllValues().get(0).payload(); + assertEquals(commitRequest.commitId(), resp1.commitId()); + assertEquals(Errors.EPOCH_MISMATCH, resp1.code()); + assertEquals(List.of(new WorkerOffset(1, 1, -1)), resp1.nextOffsets()); + + CommitResponse resp2 = ac1.getAllValues().get(1).payload(); + assertEquals(commitRequest.commitId(), resp2.commitId()); + assertEquals(Errors.NONE, resp2.code()); + assertEquals(List.of(new WorkerOffset(2, 3, 33)), resp2.nextOffsets()); + assertTrue(writers.get(0).isCompleted()); + assertEquals(3, writers.get(0).records.size()); + for (int i = 0; i < 3; i++) { + assertEquals(30 + i, writers.get(0).records.get(i).offset()); + } + } + + // 2. await commit response timeout, resend a new commit request + for (int i = 0; i < 2; i++) { + CommitRequest commitRequest2 = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(0, 1, 10), + new WorkerOffset(1, 2, 20), + new WorkerOffset(2, 3, 30))); + AtomicReference> commitCf2 = new AtomicReference<>(); + eventLoop.submit(() -> commitCf2.set(topicPartitionsWorker.commit(commitRequest2))).get(); + commitCf2.get().get(); + // expect no more read + verify(p2.log().get(), times(1)).readAsync(anyLong(), anyInt(), any(), anyBoolean()); + ArgumentCaptor ac3 = ArgumentCaptor.forClass(Event.class); + verify(channel, times((i + 1) * 2 + 2)).asyncSend(eq(TOPIC), ac3.capture()); + CommitResponse resp3 = ac3.getValue().payload(); + assertEquals(commitRequest2.commitId(), resp3.commitId()); + assertEquals(Errors.MORE_DATA, resp3.code()); + assertEquals(List.of(new WorkerOffset(2, 3, 33)), resp3.nextOffsets()); + assertEquals(1, topicPartitionsWorker.writers.size()); + assertTrue(topicPartitionsWorker.writers.get(0).isCompleted()); + } + + // 3. move to the next commit + { + CommitRequest commitRequest3 = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(0, 1, 10), + new WorkerOffset(1, 2, 20), + new WorkerOffset(2, 3, 33))); + AtomicReference> commitCf3 = new AtomicReference<>(); + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, newMemoryRecord(33, 10)); + when(p2.log().get().readAsync(eq(33L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().highWatermark()).thenReturn(43L); + eventLoop.submit(() -> commitCf3.set(topicPartitionsWorker.commit(commitRequest3))).get(); + commitCf3.get().get(); + verify(p2.log().get(), times(2)).readAsync(anyLong(), anyInt(), any(), anyBoolean()); + ArgumentCaptor ac4 = ArgumentCaptor.forClass(Event.class); + verify(channel, times(8)).asyncSend(eq(TOPIC), ac4.capture()); + CommitResponse resp4 = ac4.getValue().payload(); + assertEquals(commitRequest3.commitId(), resp4.commitId()); + assertEquals(Errors.NONE, resp4.code()); + assertEquals(List.of(new WorkerOffset(2, 3, 43)), resp4.nextOffsets()); + assertEquals(1, topicPartitionsWorker.writers.size()); + MemoryWriter writer = (MemoryWriter) topicPartitionsWorker.writers.get(0); + assertEquals(10, writer.records.size()); + for (int j = 0; j < 10; j++) { + assertEquals(33 + j, writer.records.get(j).offset()); + } + } + eventLoop.submit(() -> assertEquals(4, commitLimiter.availablePermits())).get(); + } + + @Test + public void testAdvanceSync() throws Exception { + Partition p2 = partition(2, 3); + + { + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, newMemoryRecord(30, 3)); + when(p2.log().get().readAsync(eq(30L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().highWatermark()).thenReturn(33L); + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(1, 2, 20), + new WorkerOffset(2, 3, 30) + )); + AtomicReference> commitCf = new AtomicReference<>(); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p2); + commitCf.set(topicPartitionsWorker.commit(commitRequest)); + }).get(); + commitCf.get().get(); + } + + { + avgRecordRecord = 512; + MemoryRecords records = newMemoryRecord(33, 4, 512); + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, records); + when(p2.log().get().readAsync(eq(33L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().readAsync(eq(35L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().readAsync(eq(37L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(FetchDataInfo.empty(37L))); + when(p2.log().get().highWatermark()).thenReturn(37L); + + topicPartitionsWorker.advanceSync.onAppend(new TopicPartition(TOPIC, 2), records); + + eventLoop.submit(() -> topicPartitionsWorker.run()).get(); + + for (; ; ) { + if (topicPartitionsWorker.status == TopicPartitionsWorker.Status.IDLE) { + break; + } + Threads.sleep(10); + } + assertEquals(3, writers.size()); + assertWriter(writers.get(1), 33, 35); + assertWriter(writers.get(2), 35, 37); + assertEquals(4, commitLimiter.availablePermits()); + } + + // new commit request + { + avgRecordRecord = 1; + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(2, 3, 33))); + AtomicReference> commitCf = new AtomicReference<>(); + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, newMemoryRecord(37, 10)); + when(p2.log().get().readAsync(eq(37L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().highWatermark()).thenReturn(47L); + eventLoop.submit(() -> commitCf.set(topicPartitionsWorker.commit(commitRequest))).get(); + commitCf.get().get(); + ArgumentCaptor ac = ArgumentCaptor.forClass(Event.class); + verify(channel, times(2)).asyncSend(eq(TOPIC), ac.capture()); + CommitResponse resp = ac.getValue().payload(); + assertEquals(commitRequest.commitId(), resp.commitId()); + assertEquals(Errors.NONE, resp.code()); + assertEquals(List.of(new WorkerOffset(2, 3, 47)), resp.nextOffsets()); + assertEquals(3, topicPartitionsWorker.writers.size()); + MemoryWriter writer = (MemoryWriter) topicPartitionsWorker.writers.get(2); + assertEquals(10, writer.records.size()); + assertWriter((MemoryWriter) topicPartitionsWorker.writers.get(2), 37, 47); + } + } + + @Test + public void testPartitionChange() throws Exception { + Partition p2 = partition(2, 3); + + { + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, newMemoryRecord(30, 3)); + when(p2.log().get().readAsync(eq(30L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().highWatermark()).thenReturn(33L); + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(1, 1, 10), + new WorkerOffset(2, 3, 30) + )); + AtomicReference> commitCf = new AtomicReference<>(); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p2); + commitCf.set(topicPartitionsWorker.commit(commitRequest)); + }).get(); + commitCf.get().get(); + } + + Partition p1 = partition(1, 1); + { + when(p1.log().get().readAsync(eq(10L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(new FetchDataInfo(null, newMemoryRecord(10, 5)))); + when(p1.log().get().highWatermark()).thenReturn(15L); + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(1, 1, 10), + new WorkerOffset(2, 3, 30) + )); + AtomicReference> commitCf = new AtomicReference<>(); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p1); + commitCf.set(topicPartitionsWorker.commit(commitRequest)); + }).get(); + commitCf.get().get(); + ArgumentCaptor ac = ArgumentCaptor.forClass(Event.class); + verify(channel, times(2)).asyncSend(eq(TOPIC), ac.capture()); + CommitResponse resp = ac.getValue().payload(); + assertEquals(commitRequest.commitId(), resp.commitId()); + assertEquals(Errors.NONE, resp.code()); + assertEquals(List.of(new WorkerOffset(1, 1, 15), new WorkerOffset(2, 3, 33)), resp.nextOffsets()); + assertEquals(1, topicPartitionsWorker.writers.size()); + verify(p1.log().get(), times(1)).readAsync(anyLong(), anyInt(), any(), anyBoolean()); + // reset the writer and re-read + verify(p2.log().get(), times(2)).readAsync(anyLong(), anyInt(), any(), anyBoolean()); + } + } + + @Test + public void testFlushFail() throws Exception { + MemoryWriter writer = spy(new MemoryWriter(config)); + when(writerFactory.newWriter()).thenReturn(writer); + doReturn(CompletableFuture.failedFuture(new IOException("test io exception"))).when(writer).flush(any(), any(), any()); + + Partition p2 = partition(2, 3); + + { + FetchDataInfo fetchDataInfo = new FetchDataInfo(null, newMemoryRecord(30, 3)); + when(p2.log().get().readAsync(eq(30L), anyInt(), any(), eq(true))).thenReturn(CompletableFuture.completedFuture(fetchDataInfo)); + when(p2.log().get().highWatermark()).thenReturn(33L); + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(2, 3, 30) + )); + AtomicReference> commitCf = new AtomicReference<>(); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p2); + commitCf.set(topicPartitionsWorker.commit(commitRequest)); + }).get(); + commitCf.get().get(); + ArgumentCaptor ac = ArgumentCaptor.forClass(Event.class); + verify(channel, times(1)).asyncSend(eq(TOPIC), ac.capture()); + CommitResponse resp = ac.getValue().payload(); + assertEquals(commitRequest.commitId(), resp.commitId()); + assertEquals(Errors.MORE_DATA, resp.code()); + assertEquals(List.of(new WorkerOffset(2, 3, 30)), resp.nextOffsets()); + } + + MemoryWriter writer2 = new MemoryWriter(config); + when(writerFactory.newWriter()).thenReturn(writer2); + { + CommitRequest commitRequest = new CommitRequest(UUID.randomUUID(), TOPIC, List.of( + new WorkerOffset(2, 3, 30) + )); + AtomicReference> commitCf = new AtomicReference<>(); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p2); + commitCf.set(topicPartitionsWorker.commit(commitRequest)); + }).get(); + commitCf.get().get(); + ArgumentCaptor ac = ArgumentCaptor.forClass(Event.class); + verify(channel, times(2)).asyncSend(eq(TOPIC), ac.capture()); + CommitResponse resp = ac.getValue().payload(); + assertEquals(commitRequest.commitId(), resp.commitId()); + assertEquals(Errors.NONE, resp.code()); + assertEquals(List.of(new WorkerOffset(2, 3, 33)), resp.nextOffsets()); + } + + } + + @Test + public void testSyncTask() throws ExecutionException, InterruptedException { + Partition p1 = partition(1, 2); + Partition p2 = partition(2, 3); + eventLoop.submit(() -> { + topicPartitionsWorker.add(p1); + topicPartitionsWorker.add(p2); + + avgRecordRecord = 128; + when(p1.log().get().highWatermark()).thenReturn(50L); + when(p2.log().get().highWatermark()).thenReturn(100L); + + SyncTask syncTask = topicPartitionsWorker.newSyncTask("test", Map.of(1, 10L, 2, 20L), 0); + syncTask.plan(); + assertTrue(syncTask.hasMoreData()); + List tasks = syncTask.microSyncBatchTasks; + assertEquals(2, tasks.size()); + assertEquals(List.of(new OffsetBound(p2, 20L, 28L)), tasks.get(0).offsetBounds()); + assertEquals(List.of(new OffsetBound(p1, 10L, 18L)), tasks.get(1).offsetBounds()); + }).get(); + } + + @Test + public void testMicroSyncBatchTask() throws ExecutionException, InterruptedException { + Partition p1 = partition(1, 2); + Partition p2 = partition(2, 3); + { + EventLoops.EventLoopRef eventLoopRef = new EventLoops.EventLoopRef(new EventLoops.EventLoopWrapper(eventLoop)); + PartitionWriteTaskContext ctx = new PartitionWriteTaskContext(new MemoryWriter(config), eventLoopRef, flushExecutor, config, 0); + MicroSyncBatchTask task = spy(new MicroSyncBatchTask("test", ctx)); + task.addPartition(new OffsetBound(p1, 100, 200)); + task.addPartition(new OffsetBound(p2, 10, 50)); + task.startOffsets(Map.of(1, 110L, 2, 30L)); + task.endOffsets(100, 1); + assertEquals(List.of(new OffsetBound(p1, 110L, 190L), new OffsetBound(p2, 30L, 50L)), task.offsetBounds()); + assertTrue(task.hasMoreData()); + task.run().get(); + verify(task, times(1)).runPartitionWriteTask(eq(p1), eq(110L), eq(190L)); + verify(task, times(1)).runPartitionWriteTask(eq(p2), eq(30L), eq(50L)); + assertEquals(0, eventLoopRef.inflight.get()); + } + { + EventLoops.EventLoopRef eventLoopRef = new EventLoops.EventLoopRef(new EventLoops.EventLoopWrapper(eventLoop)); + PartitionWriteTaskContext ctx = new PartitionWriteTaskContext(new MemoryWriter(config), eventLoopRef, flushExecutor, config, 0); + MicroSyncBatchTask task = spy(new MicroSyncBatchTask("test", ctx)); + task.addPartition(new OffsetBound(p1, 100, 200)); + task.addPartition(new OffsetBound(p2, 10, 50)); + task.startOffsets(Map.of(1, 110L, 2, 30L)); + task.endOffsets(200, 1); + assertEquals(List.of(new OffsetBound(p1, 110L, 200L), new OffsetBound(p2, 30L, 50L)), task.offsetBounds()); + assertFalse(task.hasMoreData()); + } + } + + private static void assertWriter(MemoryWriter writer, long start, long end) { + assertTrue(writer.isCompleted()); + assertEquals(end - start, writer.records.size()); + for (int i = 0; i < writer.records.size(); i++) { + assertEquals(start + i, writer.records.get(i).offset()); + } + } + + class TopicPartitionsWorker extends kafka.automq.table.worker.TopicPartitionsWorker { + + public TopicPartitionsWorker(String topic, WorkerConfig config, WriterFactory writerFactory, Channel channel, + EventLoop eventLoop, EventLoops eventLoops, ExecutorService flushExecutors, Semaphore commitLimiter) { + super(topic, config, writerFactory, channel, eventLoop, eventLoops, flushExecutors, commitLimiter); + } + + @Override + protected double getAvgRecordSize() { + return avgRecordRecord; + } + + @Override + protected double getDecompressedRatio() { + return 1.0; + } + } + + class MicroSyncBatchTask extends kafka.automq.table.worker.TopicPartitionsWorker.MicroSyncBatchTask { + public MicroSyncBatchTask(String logContext, PartitionWriteTaskContext ctx) { + super(logContext, ctx); + } + + @Override + protected CompletableFuture runPartitionWriteTask(Partition partition, long start, long end) { + return CompletableFuture.completedFuture(null); + } + } + +} diff --git a/core/src/test/java/kafka/automq/utils/ClientUtilsTest.java b/core/src/test/java/kafka/automq/utils/ClientUtilsTest.java new file mode 100644 index 0000000000..193d1fea10 --- /dev/null +++ b/core/src/test/java/kafka/automq/utils/ClientUtilsTest.java @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.utils; + +import kafka.cluster.EndPoint; +import kafka.server.KafkaConfig; + +import org.apache.kafka.common.config.internals.BrokerSecurityConfigs; +import org.apache.kafka.common.config.types.Password; +import org.apache.kafka.common.network.ListenerName; +import org.apache.kafka.common.security.auth.SecurityProtocol; +import org.apache.kafka.network.SocketServerConfigs; +import org.apache.kafka.server.config.ReplicationConfigs; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ClientUtilsTest { + + @Test + public void clusterConfigIncludesListenerSecuritySettings() { + Properties props = baseBrokerConfig("INTERNAL", SecurityProtocol.SASL_SSL, "SCRAM-SHA-512", 19093); + String listenerLower = "internal"; + String mechanismLower = "scram-sha-512"; + props.put("listener.name." + listenerLower + ".ssl.keystore.location", "/path/keystore.jks"); + props.put("listener.name." + listenerLower + ".ssl.truststore.location", "/path/truststore.jks"); + props.put("listener.name." + listenerLower + "." + mechanismLower + ".sasl.jaas.config", "listener-jaas"); + props.put("ssl.truststore.password", "secret"); + + KafkaConfig kafkaConfig = KafkaConfig.fromProps(props); + Properties clientProps = ClientUtils.clusterClientBaseConfig(kafkaConfig); + + assertEquals("broker.local:19093", clientProps.getProperty("bootstrap.servers")); + assertEquals(SecurityProtocol.SASL_SSL.name(), clientProps.getProperty("security.protocol")); + assertEquals("SCRAM-SHA-512", clientProps.getProperty("sasl.mechanism")); + assertEquals("/path/keystore.jks", clientProps.get("ssl.keystore.location")); + assertEquals("/path/truststore.jks", clientProps.get("ssl.truststore.location")); + Object jaasConfig = clientProps.get("sasl.jaas.config"); + assertNotNull(jaasConfig); + assertEquals("listener-jaas", jaasConfig.toString()); + Object truststorePassword = clientProps.get("ssl.truststore.password"); + assertNotNull(truststorePassword); + if (truststorePassword instanceof Password) { + assertEquals("secret", ((Password) truststorePassword).value()); + } else { + assertEquals("secret", truststorePassword.toString()); + } + } + + @Test + public void setsSaslMechanismWhenAbsentInListenerConfigs() { + Properties props = baseBrokerConfig("INTERNAL", SecurityProtocol.SASL_SSL, "SCRAM-SHA-256", 19094); + KafkaConfig kafkaConfig = KafkaConfig.fromProps(props); + + Properties clientProps = ClientUtils.clusterClientBaseConfig(kafkaConfig); + + assertEquals("SCRAM-SHA-256", clientProps.getProperty("sasl.mechanism")); + } + + @Test + public void throwsWhenInterBrokerEndpointMissing() { + KafkaConfig kafkaConfig = mock(KafkaConfig.class); + ListenerName listenerName = new ListenerName("INTERNAL"); + when(kafkaConfig.interBrokerListenerName()).thenReturn(listenerName); + + List endpoints = new ArrayList<>(); + endpoints.add(new EndPoint("broker.local", 9092, new ListenerName("EXTERNAL"), SecurityProtocol.PLAINTEXT)); + when(kafkaConfig.effectiveAdvertisedBrokerListeners()).thenReturn(scalaEndpoints(endpoints)); + + assertThrows(IllegalArgumentException.class, () -> ClientUtils.clusterClientBaseConfig(kafkaConfig)); + } + + private Properties baseBrokerConfig(String listenerName, + SecurityProtocol securityProtocol, + String saslMechanism, + int port) { + Properties props = kafka.utils.TestUtils.createDummyBrokerConfig(); + String listener = listenerName.toUpperCase(Locale.ROOT); + String listenerLower = listenerName.toLowerCase(Locale.ROOT); + props.put(SocketServerConfigs.LISTENERS_CONFIG, listener + "://broker.local:" + port); + props.put(SocketServerConfigs.ADVERTISED_LISTENERS_CONFIG, listener + "://broker.local:" + port); + props.put(SocketServerConfigs.LISTENER_SECURITY_PROTOCOL_MAP_CONFIG, listener + ":" + securityProtocol.name() + ",CONTROLLER:PLAINTEXT"); + props.put(ReplicationConfigs.INTER_BROKER_LISTENER_NAME_CONFIG, listener); + props.remove(ReplicationConfigs.INTER_BROKER_SECURITY_PROTOCOL_CONFIG); + props.put(BrokerSecurityConfigs.SASL_MECHANISM_INTER_BROKER_PROTOCOL_CONFIG, saslMechanism); + props.put("listener.name." + listenerLower + ".sasl.enabled.mechanisms", saslMechanism); + return props; + } + + private scala.collection.Seq scalaEndpoints(List endpoints) { + return scala.jdk.javaapi.CollectionConverters.asScala(endpoints).toSeq(); + } +} diff --git a/core/src/test/java/kafka/automq/zerozone/DefaultClientRackProviderTest.java b/core/src/test/java/kafka/automq/zerozone/DefaultClientRackProviderTest.java new file mode 100644 index 0000000000..47644498d5 --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/DefaultClientRackProviderTest.java @@ -0,0 +1,51 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.zerozone.DefaultClientRackProvider.CIDRMatcher; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; + +@Timeout(60) +@Tag("S3Unit") +public class DefaultClientRackProviderTest { + + @Test + public void testCIDRFind() { + CIDRMatcher matcher = new CIDRMatcher("us-east-1a@10.0.0.0/19,10.0.32.0/19<>us-east-1b@10.0.64.0/19<>us-east-1c@10.0.96.0/19"); + assertEquals("us-east-1a", matcher.find("10.0.31.233").zone()); + assertEquals("10.0.0.0/19", matcher.find("10.0.31.233").cidr()); + + assertEquals("us-east-1a", matcher.find("10.0.32.233").zone()); + assertEquals("10.0.32.0/19", matcher.find("10.0.32.233").cidr()); + + assertEquals("us-east-1b", matcher.find("10.0.65.0").zone()); + + assertEquals("us-east-1c", matcher.find("10.0.97.0").zone()); + + assertNull(matcher.find("10.0.128.0")); + } + +} diff --git a/core/src/test/java/kafka/automq/zerozone/LinkRecordTest.java b/core/src/test/java/kafka/automq/zerozone/LinkRecordTest.java new file mode 100644 index 0000000000..55ddaf1dd1 --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/LinkRecordTest.java @@ -0,0 +1,54 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.SimpleRecord; + +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; + +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; + +import io.netty.buffer.ByteBuf; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class LinkRecordTest { + + @Test + public void testEncodeDecode() { + SimpleRecord record = new SimpleRecord(ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))); + MemoryRecords memoryRecords = MemoryRecords.withRecords(Compression.NONE, record); + + ByteBuf linkRecordBuf = LinkRecord.encode( + ChannelOffset.of( + (short) 1, (short) 2, 3, 4, + DefaultRecordOffset.of(5, 6, 7).buffer() + ), + memoryRecords); + + assertEquals(7, LinkRecord.decodedSize(linkRecordBuf)); + } + +} diff --git a/core/src/test/java/kafka/automq/zerozone/ProxyNodeMappingTest.java b/core/src/test/java/kafka/automq/zerozone/ProxyNodeMappingTest.java new file mode 100644 index 0000000000..4038cfd751 --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/ProxyNodeMappingTest.java @@ -0,0 +1,182 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.interceptor.ClientIdMetadata; +import kafka.automq.zerozone.ProxyNodeMapping.ProxyNode; +import kafka.server.MetadataCache; + +import org.apache.kafka.common.Endpoint; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.security.auth.SecurityProtocol; +import org.apache.kafka.metadata.BrokerRegistration; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Timeout(60) +@Tag("S3Unit") +public class ProxyNodeMappingTest { + static final String LISTENER_NAME = "BROKER"; + MetadataCache metadataCache; + ProxyNodeMapping proxyNodeMapping; + + @BeforeEach + public void setup() { + metadataCache = mock(MetadataCache.class); + proxyNodeMapping = new ProxyNodeMapping(new Node(1, "127.0.0.1", 9092), "az1", LISTENER_NAME, metadataCache); + + Map> main2proxyByRack = new HashMap<>(); + HashMap az1 = new HashMap<>(); + az1.put(3, brokerRegistration(2, "az1", "127.0.0.2", 9092)); + az1.put(4, brokerRegistration(1, "az1", "127.0.0.1", 9092)); + main2proxyByRack.put("az1", az1); + HashMap az2 = new HashMap<>(); + az2.put(1, brokerRegistration(3, "az2", "127.0.0.3", 9092)); + az2.put(2, brokerRegistration(3, "az2", "127.0.0.3", 9092)); + main2proxyByRack.put("az2", az2); + + proxyNodeMapping.main2proxyByRack = main2proxyByRack; + } + + @Test + public void testGetRouteOutNode() { + Node node; + // case1 + when(metadataCache.getPartitionLeaderNode(eq("TP"), eq(1))) + .thenReturn(brokerRegistration(3, "az2", "127.0.0.3", 9092)); + node = proxyNodeMapping.getRouteOutNode("TP", 1, ClientIdMetadata.of("automq_az=az1")); + assertEquals(-1, node.id()); + + // case2 + when(metadataCache.getPartitionLeaderNode(eq("TP"), eq(1))) + .thenReturn(brokerRegistration(2, "az1", "127.0.0.2", 9092)); + node = proxyNodeMapping.getRouteOutNode("TP", 1, ClientIdMetadata.of("automq_az=az1")); + assertEquals(-1, node.id()); + + // case3 + when(metadataCache.getPartitionLeaderNode(eq("TP"), eq(1))) + .thenReturn(brokerRegistration(1, "az1", "127.0.0.1", 9092)); + node = proxyNodeMapping.getRouteOutNode("TP", 1, ClientIdMetadata.of("automq_az=az1")); + assertEquals(1, node.id()); + + // case4 + when(metadataCache.getPartitionLeaderNode(eq("TP"), eq(1))) + .thenReturn(brokerRegistration(1, "az1", "127.0.0.1", 9092)); + node = proxyNodeMapping.getRouteOutNode("TP", 1, ClientIdMetadata.of("automq_az=az3")); + assertEquals(1, node.id()); + + // case5 + when(metadataCache.getPartitionLeaderNode(eq("TP"), eq(1))) + .thenReturn(brokerRegistration(3, "az2", "127.0.0.3", 9092)); + node = proxyNodeMapping.getRouteOutNode("TP", 1, ClientIdMetadata.of("automq_az=az3")); + assertEquals(-1, node.id()); + + // case6 + when(metadataCache.getPartitionLeaderNode(eq("TP"), eq(1))) + .thenReturn(brokerRegistration(4, "az2", "127.0.0.4", 9092)); + node = proxyNodeMapping.getRouteOutNode("TP", 1, ClientIdMetadata.of("automq_az=az1")); + assertEquals(4, node.id()); + } + + @Test + public void testGetLeaderNode() { + Node node; + // case1 + when(metadataCache.getNode(eq(3))) + .thenReturn(brokerRegistration(3, "az2", "127.0.0.3", 9092)); + node = proxyNodeMapping.getLeaderNode(3, ClientIdMetadata.of(""), LISTENER_NAME).get(); + assertEquals(3, node.id()); + + // case2 + when(metadataCache.getNode(eq(3))) + .thenReturn(brokerRegistration(3, "az2", "127.0.0.3", 9092)); + node = proxyNodeMapping.getLeaderNode(3, ClientIdMetadata.of("automq_az=az1"), LISTENER_NAME).get(); + assertEquals(2, node.id()); + + // case3 + when(metadataCache.getNode(eq(3))) + .thenReturn(brokerRegistration(3, "az2", "127.0.0.3", 9092)); + node = proxyNodeMapping.getLeaderNode(3, ClientIdMetadata.of("automq_az=az3"), LISTENER_NAME).get(); + assertEquals(3, node.id()); + + // case4 + when(metadataCache.getNode(eq(2))) + .thenReturn(brokerRegistration(2, "az1", "127.0.0.2", 9092)); + node = proxyNodeMapping.getLeaderNode(2, ClientIdMetadata.of("automq_az=az1"), LISTENER_NAME).get(); + assertEquals(2, node.id()); + } + + @Test + public void testCalMain2proxyByRack() { + Map> main2proxyByRack = new HashMap<>(); + List az1 = new ArrayList<>(); + az1.add(brokerRegistration(1, "az1", "127.0.0.1", 9092)); + az1.add(brokerRegistration(2, "az1", "127.0.0.2", 9092)); + main2proxyByRack.put("az1", az1); + List az2 = new ArrayList<>(); + az2.add(brokerRegistration(3, "az2", "127.0.0.3", 9092)); + az2.add(brokerRegistration(4, "az2", "127.0.0.4", 9092)); + az2.add(brokerRegistration(5, "az2", "127.0.0.5", 9092)); + main2proxyByRack.put("az2", az2); + + Map> rst = ProxyNodeMapping.calMain2proxyByRack(main2proxyByRack); + assertEquals(2, rst.size()); + assertEquals(3, rst.get("az1").size()); + assertEquals(List.of(3, 4, 5), rst.get("az1").keySet().stream().sorted().toList()); + assertEquals(2, rst.get("az2").size()); + assertEquals(List.of(1, 2), rst.get("az2").keySet().stream().sorted().toList()); + } + + @Test + public void testTryFreeController() { + List proxyNodes = new ArrayList<>(); + ProxyNode node1 = new ProxyNode(brokerRegistration(1, "az1", "127.0.0.1", 9092)); + node1.mainNodeIds.addAll(List.of(1000, 1001)); + proxyNodes.add(node1); + ProxyNode node2 = new ProxyNode(brokerRegistration(2000, "az1", "127.0.0.1", 9092)); + node2.mainNodeIds.add(0); + proxyNodes.add(node2); + ProxyNodeMapping.tryFreeController(proxyNodes, 2); + assertEquals(List.of(0), node1.mainNodeIds); + assertEquals(List.of(1001, 1000), node2.mainNodeIds); + } + + private static BrokerRegistration brokerRegistration(int id, String rack, String host, int port) { + return new BrokerRegistration.Builder() + .setId(id).setRack(Optional.of(rack)) + .setListeners(List.of(new Endpoint(LISTENER_NAME, SecurityProtocol.PLAINTEXT, host, port))) + .build(); + } + +} diff --git a/core/src/test/java/kafka/automq/zerozone/RouterRecordV2Test.java b/core/src/test/java/kafka/automq/zerozone/RouterRecordV2Test.java new file mode 100644 index 0000000000..a7377cefbd --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/RouterRecordV2Test.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.List; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +@Timeout(60) +@Tag("S3Unit") +public class RouterRecordV2Test { + + @Test + public void test_encodeDecode() { + int nodeId = 233; + ByteBuf channelOffset1 = Unpooled.buffer(); + channelOffset1.writeInt(1); + ByteBuf channelOffset2 = Unpooled.buffer(); + channelOffset2.writeInt(2); + + RouterRecordV2 record = RouterRecordV2.decode(new RouterRecordV2(nodeId, List.of(channelOffset1, channelOffset2)).encode()); + + assertEquals(nodeId, record.nodeId()); + assertEquals(2, record.channelOffsets().size()); + assertEquals(4, record.channelOffsets().get(0).readableBytes()); + assertEquals(1, record.channelOffsets().get(0).readInt()); + assertEquals(4, record.channelOffsets().get(1).readableBytes()); + assertEquals(2, record.channelOffsets().get(1).readInt()); + } + +} diff --git a/core/src/test/java/kafka/automq/zerozone/SnapshotReadPartitionsManagerTest.java b/core/src/test/java/kafka/automq/zerozone/SnapshotReadPartitionsManagerTest.java new file mode 100644 index 0000000000..7afd496d8f --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/SnapshotReadPartitionsManagerTest.java @@ -0,0 +1,250 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import kafka.automq.partition.snapshot.SnapshotOperation; +import kafka.automq.zerozone.SnapshotReadPartitionsManager.OperationBatch; +import kafka.automq.zerozone.SnapshotReadPartitionsManager.Subscriber; +import kafka.cluster.Partition; +import kafka.cluster.PartitionSnapshot; +import kafka.server.KafkaConfig; +import kafka.server.MetadataCache; +import kafka.server.streamaspect.ElasticReplicaManager; + +import org.apache.kafka.common.Endpoint; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.TopicIdPartition; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AutomqGetPartitionSnapshotResponseData; +import org.apache.kafka.common.network.ListenerName; +import org.apache.kafka.common.security.auth.SecurityProtocol; +import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.metadata.BrokerRegistration; +import org.apache.kafka.server.common.automq.AutoMQVersion; + +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.WriteAheadLog; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.OptionalLong; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Timeout(60) +@Tag("S3Unit") +public class SnapshotReadPartitionsManagerTest { + KafkaConfig config; + MockTime time; + + ElasticReplicaManager replicaManager; + Map partitions; + + MetadataCache metadataCache; + Map stream2offset; + Map topicId2name; + + AsyncSender asyncSender; + + SnapshotReadPartitionsManager manager; + + String topicName; + Uuid topicId; + + BrokerRegistration broker1 = new BrokerRegistration.Builder().setId(1).setListeners(List.of(new Endpoint("BROKER", SecurityProtocol.PLAINTEXT, "127.0.0.1", 9092))).build(); + BrokerRegistration broker2 = new BrokerRegistration.Builder().setId(2).setListeners(List.of(new Endpoint("BROKER", SecurityProtocol.PLAINTEXT, "127.0.0.2", 9092))).build(); + + @BeforeEach + public void setup() { + config = mock(KafkaConfig.class); + when(config.nodeId()).thenReturn(1); + when(config.interBrokerListenerName()).thenReturn(ListenerName.normalised("BROKER")); + + time = new MockTime(); + + replicaManager = mock(ElasticReplicaManager.class); + partitions = new HashMap<>(); + doAnswer(args -> { + TopicPartition tp = args.getArgument(0); + return partitions.compute(tp, args.getArgument(1)); + }).when(replicaManager).computeSnapshotReadPartition(any(), any()); + doAnswer(args -> mock(Partition.class)).when(replicaManager).newSnapshotReadPartition(any()); + + metadataCache = mock(MetadataCache.class); + stream2offset = new HashMap<>(); + doAnswer(args -> { + Long offset = stream2offset.get((Long) args.getArgument(0)); + if (offset == null) { + return OptionalLong.empty(); + } else { + return OptionalLong.of(offset); + } + }).when(metadataCache).getStreamEndOffset(anyLong()); + topicId2name = new HashMap<>(); + topicName = "topic1"; + topicId = Uuid.randomUuid(); + topicId2name.put(topicId, topicName); + when(metadataCache.topicIdsToNames()).thenReturn(topicId2name); + + asyncSender = mock(AsyncSender.class); + + Replayer replayer = new Replayer() { + @Override + public CompletableFuture replay(List objects) { + return CompletableFuture.completedFuture(null); + } + + @Override + public CompletableFuture replay(WriteAheadLog confirmWAL, RecordOffset startOffset, RecordOffset endOffset, List walRecords) { + return CompletableFuture.completedFuture(null); + } + }; + + manager = new SnapshotReadPartitionsManager(config, time, null, replicaManager, metadataCache, replayer, asyncSender); + } + + @Test + public void testSubscriber_zerozonev1() throws Exception { + Node node = new Node(2, "127.0.0.1", 9092); + SubscriberRequester requester = mock(SubscriberRequester.class); + SubscriberReplayer dataLoader = mock(SubscriberReplayer.class); + when(dataLoader.replayWal()).thenReturn(CompletableFuture.completedFuture(null)); + + Subscriber subscriber = manager.newSubscriber(node, AutoMQVersion.V3, requester, dataLoader); + + OperationBatch operationBatch = new OperationBatch(); + operationBatch.operations.add(snapshotWithOperation(1, Map.of(3L, 3L), SnapshotOperation.ADD)); + operationBatch.operations.add(snapshotWithOperation(2, Map.of(4L, 6L), SnapshotOperation.ADD)); + subscriber.onNewOperationBatch(operationBatch); + + subscriber.run(); + verify(dataLoader, timeout(1000L).times(1)).replayWal(); + awaitEventLoopClear(); + + assertEquals(2, subscriber.partitions.size()); + verify(partitions.get(new TopicPartition(topicName, 1)), times(1)).snapshot(any()); + verify(partitions.get(new TopicPartition(topicName, 2)), times(1)).snapshot(any()); + } + + @Test + public void testSubscriber_zerozonev0() throws ExecutionException, InterruptedException { + Node node = new Node(2, "127.0.0.1", 9092); + SubscriberRequester requester = mock(SubscriberRequester.class); + SubscriberReplayer dataLoader = mock(SubscriberReplayer.class); + when(dataLoader.relayObject()).thenReturn(CompletableFuture.completedFuture(null)); + Subscriber subscriber = manager.newSubscriber(node, AutoMQVersion.V2, requester, dataLoader); + stream2offset.put(3L, 0L); + stream2offset.put(4L, 0L); + + OperationBatch operationBatch = new OperationBatch(); + operationBatch.operations.add(snapshotWithOperation(1, Map.of(3L, 3L), SnapshotOperation.ADD)); + operationBatch.operations.add(snapshotWithOperation(2, Map.of(4L, 6L), SnapshotOperation.ADD)); + subscriber.onNewOperationBatch(operationBatch); + + // metadata unready. + subscriber.tryReplay(); + assertEquals(1, subscriber.waitingMetadataReadyQueue.size()); + assertEquals(-1, operationBatch.readyIndex); + + // metadata ready. + stream2offset.put(3L, 3L); + stream2offset.put(4L, 6L); + subscriber.tryReplay(); + awaitEventLoopClear(); + assertEquals(2, subscriber.partitions.size()); + verify(partitions.get(new TopicPartition(topicName, 1)), times(1)).snapshot(any()); + verify(partitions.get(new TopicPartition(topicName, 2)), times(1)).snapshot(any()); + assertEquals(0, subscriber.waitingMetadataReadyQueue.size()); + assertEquals(0, subscriber.waitingDataLoadedQueue.size()); + assertEquals(0, subscriber.snapshotWithOperations.size()); + + // partition2 append new records + operationBatch = new OperationBatch(); + operationBatch.operations.add(snapshotWithOperation(2, Map.of(4L, 10L), SnapshotOperation.PATCH)); + subscriber.onNewOperationBatch(operationBatch); + stream2offset.put(4L, 10L); + subscriber.tryReplay(); + awaitEventLoopClear(); + verify(partitions.get(new TopicPartition(topicName, 1)), times(1)).snapshot(any()); + verify(partitions.get(new TopicPartition(topicName, 2)), times(2)).snapshot(any()); + + // partition2 remove + operationBatch = new OperationBatch(); + operationBatch.operations.add(snapshotWithOperation(2, Map.of(4L, 10L), SnapshotOperation.REMOVE)); + subscriber.onNewOperationBatch(operationBatch); + subscriber.tryReplay(); + awaitEventLoopClear(); + assertEquals(1, partitions.size()); + assertEquals(1, subscriber.partitions.size()); + + } + + private SnapshotWithOperation snapshotWithOperation(int partitionIndex, Map offsets, + SnapshotOperation operation) { + // TODO: fix the test + PartitionSnapshot snapshot = new PartitionSnapshot(0, null, null, null, offsets, null, CompletableFuture.completedFuture(null)); + return new SnapshotWithOperation(new TopicIdPartition(topicId, partitionIndex, topicName), snapshot, operation); + } + + private void awaitEventLoopClear() throws ExecutionException, InterruptedException { + manager.eventLoop.submit(() -> { + }).get(); + } + + static AutomqGetPartitionSnapshotResponseData snapshotResponse(Uuid topicId, int partitionId, int leaderEpoch, + long endOffset, SnapshotOperation ops) { + AutomqGetPartitionSnapshotResponseData data = new AutomqGetPartitionSnapshotResponseData(); + AutomqGetPartitionSnapshotResponseData.TopicCollection topics = new AutomqGetPartitionSnapshotResponseData.TopicCollection(); + AutomqGetPartitionSnapshotResponseData.Topic topic = new AutomqGetPartitionSnapshotResponseData.Topic(); + topic.setTopicId(topicId); + AutomqGetPartitionSnapshotResponseData.PartitionSnapshot snapshot = new AutomqGetPartitionSnapshotResponseData.PartitionSnapshot(); + snapshot.setPartitionIndex(partitionId); + snapshot.setOperation(ops.code()); + snapshot.setLeaderEpoch(leaderEpoch); + snapshot.setLogMetadata(new AutomqGetPartitionSnapshotResponseData.LogMetadata()); + snapshot.setFirstUnstableOffset(null); + snapshot.setLogEndOffset(new AutomqGetPartitionSnapshotResponseData.LogOffsetMetadata().setMessageOffset(endOffset).setRelativePositionInSegment(1)); + snapshot.setStreamMetadata(List.of(new AutomqGetPartitionSnapshotResponseData.StreamMetadata().setStreamId(partitionId).setEndOffset(endOffset))); + topic.setPartitions(List.of(snapshot)); + topics.add(topic); + data.setTopics(topics); + return data; + } + +} diff --git a/core/src/test/java/kafka/automq/zerozone/ZoneRouterPackTest.java b/core/src/test/java/kafka/automq/zerozone/ZoneRouterPackTest.java new file mode 100644 index 0000000000..e91ca1cfa0 --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/ZoneRouterPackTest.java @@ -0,0 +1,67 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.compress.Compression; +import org.apache.kafka.common.message.ProduceRequestData; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.SimpleRecord; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; + +import io.netty.buffer.ByteBuf; + +@Timeout(60) +@Tag("S3Unit") +public class ZoneRouterPackTest { + + @Test + public void testDataBlockCodec() { + for (short version : new short[] {3, 7, 11}) { + List produceRequests = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + ProduceRequestData requestData = new ProduceRequestData(); + requestData.setTransactionalId("trans" + i); + requestData.setAcks((short) -1); + ProduceRequestData.TopicProduceDataCollection produceData = new ProduceRequestData.TopicProduceDataCollection(); + ProduceRequestData.PartitionProduceData partitionProduceData = new ProduceRequestData.PartitionProduceData(); + partitionProduceData.setIndex(i); + partitionProduceData.setRecords(MemoryRecords.withRecords(Compression.NONE, new SimpleRecord(("simplerecord" + i).getBytes(StandardCharsets.UTF_8)))); + produceData.add( + new ProduceRequestData.TopicProduceData() + .setName("topic") + .setPartitionData(List.of(partitionProduceData))); + requestData.setTopicData(produceData); + produceRequests.add(new ZoneRouterProduceRequest(version, new ZoneRouterProduceRequest.Flag().internalTopicsAllowed(true).value(), requestData)); + } + ByteBuf buf = ZoneRouterPackWriter.encodeDataBlock(produceRequests); + List decoded = ZoneRouterPackReader.decodeDataBlock(buf); + Assertions.assertEquals(produceRequests, decoded); + } + } + +} diff --git a/core/src/test/java/kafka/automq/zerozone/ZoneRouterResponseCodecTest.java b/core/src/test/java/kafka/automq/zerozone/ZoneRouterResponseCodecTest.java new file mode 100644 index 0000000000..1e382fb75e --- /dev/null +++ b/core/src/test/java/kafka/automq/zerozone/ZoneRouterResponseCodecTest.java @@ -0,0 +1,49 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.automq.zerozone; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.message.ProduceResponseData; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.ProduceResponse; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +@Timeout(60) +@Tag("S3Unit") +public class ZoneRouterResponseCodecTest { + + @Test + public void testCodec() { + Map map = new HashMap<>(); + map.put(new TopicPartition("test", 1), new ProduceResponse.PartitionResponse(Errors.UNKNOWN_LEADER_EPOCH)); + ProduceResponseData data = new ProduceResponse(map).data(); + ProduceResponseData decoded = ZoneRouterResponseCodec.decode(ZoneRouterResponseCodec.encode(data)); + assertEquals(data, decoded); + } + +} diff --git a/core/src/test/java/kafka/log/stream/s3/objects/ObjectAttributesTest.java b/core/src/test/java/kafka/log/stream/s3/objects/ObjectAttributesTest.java index 5bb85274b5..5d58c9cc8b 100644 --- a/core/src/test/java/kafka/log/stream/s3/objects/ObjectAttributesTest.java +++ b/core/src/test/java/kafka/log/stream/s3/objects/ObjectAttributesTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.stream.s3.objects; diff --git a/core/src/test/java/kafka/log/stream/s3/telemetry/exporter/MetricsExporterURITest.java b/core/src/test/java/kafka/log/stream/s3/telemetry/exporter/MetricsExporterURITest.java deleted file mode 100644 index de8fedee2c..0000000000 --- a/core/src/test/java/kafka/log/stream/s3/telemetry/exporter/MetricsExporterURITest.java +++ /dev/null @@ -1,193 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.exporter; - -import kafka.automq.AutoMQConfig; -import kafka.server.KafkaConfig; - -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; -import org.mockito.Mockito; - -public class MetricsExporterURITest { - - @Test - public void testsBackwardCompatibility() { - String clusterId = "test_cluster"; - - KafkaConfig kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getBoolean(AutoMQConfig.S3_METRICS_ENABLE_CONFIG)).thenReturn(false); - AutoMQConfig automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - MetricsExporterURI uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNull(uri); - - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.nodeId()).thenReturn(1); - Mockito.when(kafkaConfig.getBoolean(AutoMQConfig.S3_METRICS_ENABLE_CONFIG)).thenReturn(true); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_TYPE_CONFIG)).thenReturn("otlp,prometheus"); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_ENDPOINT_CONFIG)).thenReturn("http://localhost:4318"); - Mockito.when(kafkaConfig.getBoolean(AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_COMPRESSION_ENABLE_CONFIG)).thenReturn(true); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_EXPORTER_OTLP_PROTOCOL_CONFIG)).thenReturn("http"); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_METRICS_EXPORTER_PROM_HOST_CONFIG)).thenReturn("127.0.0.1"); - Mockito.when(kafkaConfig.getInt(AutoMQConfig.S3_METRICS_EXPORTER_PROM_PORT_CONFIG)).thenReturn(9999); - Mockito.when(kafkaConfig.getBoolean(AutoMQConfig.S3_TELEMETRY_OPS_ENABLED_CONFIG)).thenReturn(true); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_OPS_BUCKETS_CONFIG)).thenReturn("0@s3://bucket0?region=us-west-1"); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - Mockito.when(kafkaConfig.s3ExporterReportIntervalMs()).thenReturn(1000); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertEquals(3, uri.metricsExporters().size()); - for (MetricsExporter metricsExporter : uri.metricsExporters()) { - if (metricsExporter instanceof OTLPMetricsExporter) { - OTLPMetricsExporter otlpExporter = (OTLPMetricsExporter) metricsExporter; - Assertions.assertEquals(1000, otlpExporter.intervalMs()); - Assertions.assertEquals("http://localhost:4318", otlpExporter.endpoint()); - Assertions.assertEquals(OTLPProtocol.HTTP, otlpExporter.protocol()); - Assertions.assertEquals(OTLPCompressionType.GZIP, otlpExporter.compression()); - } else if (metricsExporter instanceof PrometheusMetricsExporter) { - PrometheusMetricsExporter promExporter = (PrometheusMetricsExporter) metricsExporter; - Assertions.assertEquals("127.0.0.1", promExporter.host()); - Assertions.assertEquals(9999, promExporter.port()); - } else if (metricsExporter instanceof OpsMetricsExporter) { - OpsMetricsExporter opsExporter = (OpsMetricsExporter) metricsExporter; - Assertions.assertEquals(clusterId, opsExporter.clusterId()); - Assertions.assertEquals(1, opsExporter.nodeId()); - Assertions.assertEquals(1000, opsExporter.intervalMs()); - Assertions.assertEquals(1, opsExporter.opsBuckets().size()); - Assertions.assertEquals("bucket0", opsExporter.opsBuckets().get(0).bucket()); - Assertions.assertEquals("us-west-1", opsExporter.opsBuckets().get(0).region()); - } else { - Assertions.fail("Unknown exporter type"); - } - } - } - - @Test - public void testParseURIString() { - String clusterId = "test_cluster"; - // test empty exporter - KafkaConfig kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn(null); - Mockito.when(kafkaConfig.getBoolean(AutoMQConfig.S3_METRICS_ENABLE_CONFIG)).thenReturn(false); - AutoMQConfig automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - MetricsExporterURI uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNull(uri); - - // test invalid uri - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn("unknown://"); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertTrue(uri.metricsExporters().isEmpty()); - - // test invalid type - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn("unknown://?"); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertTrue(uri.metricsExporters().isEmpty()); - - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn("://?"); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertTrue(uri.metricsExporters().isEmpty()); - - // test illegal otlp config - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn("otlp://?endpoint=&protocol=grpc"); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertTrue(uri.metricsExporters().isEmpty()); - - // test illegal prometheus config - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn("prometheus://?host=&port=9999"); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertTrue(uri.metricsExporters().isEmpty()); - - // test illegal ops config - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn("ops://?"); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_OPS_BUCKETS_CONFIG)).thenReturn(""); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertTrue(uri.metricsExporters().isEmpty()); - - // test multi exporter config - kafkaConfig = Mockito.mock(KafkaConfig.class); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG)).thenReturn( - "otlp://?endpoint=http://localhost:4317&protocol=http&compression=gzip," + - "prometheus://?host=127.0.0.1&port=9999," + - "ops://?"); - Mockito.when(kafkaConfig.getString(AutoMQConfig.S3_OPS_BUCKETS_CONFIG)).thenReturn("0@s3://bucket0?region=us-west-1"); - Mockito.when(kafkaConfig.s3ExporterReportIntervalMs()).thenReturn(1000); - Mockito.when(kafkaConfig.nodeId()).thenReturn(1); - automqConfig = new AutoMQConfig(); - automqConfig.setup(kafkaConfig); - Mockito.when(kafkaConfig.automq()).thenReturn(automqConfig); - - uri = MetricsExporterURI.parse(clusterId, kafkaConfig); - Assertions.assertNotNull(uri); - Assertions.assertEquals(3, uri.metricsExporters().size()); - for (MetricsExporter metricsExporter : uri.metricsExporters()) { - if (metricsExporter instanceof OTLPMetricsExporter) { - OTLPMetricsExporter otlpExporter = (OTLPMetricsExporter) metricsExporter; - Assertions.assertEquals(1000, otlpExporter.intervalMs()); - Assertions.assertEquals("http://localhost:4317", otlpExporter.endpoint()); - Assertions.assertEquals(OTLPProtocol.HTTP, otlpExporter.protocol()); - Assertions.assertEquals(OTLPCompressionType.GZIP, otlpExporter.compression()); - Assertions.assertNotNull(metricsExporter.asMetricReader()); - } else if (metricsExporter instanceof PrometheusMetricsExporter) { - PrometheusMetricsExporter promExporter = (PrometheusMetricsExporter) metricsExporter; - Assertions.assertEquals("127.0.0.1", promExporter.host()); - Assertions.assertEquals(9999, promExporter.port()); - Assertions.assertNotNull(metricsExporter.asMetricReader()); - } else if (metricsExporter instanceof OpsMetricsExporter) { - OpsMetricsExporter opsExporter = (OpsMetricsExporter) metricsExporter; - Assertions.assertEquals(clusterId, opsExporter.clusterId()); - Assertions.assertEquals(1, opsExporter.nodeId()); - Assertions.assertEquals(1000, opsExporter.intervalMs()); - Assertions.assertEquals(1, opsExporter.opsBuckets().size()); - Assertions.assertEquals("bucket0", opsExporter.opsBuckets().get(0).bucket()); - Assertions.assertEquals("us-west-1", opsExporter.opsBuckets().get(0).region()); - } else { - Assertions.fail("Unknown exporter type"); - } - } - } -} diff --git a/core/src/test/java/kafka/log/stream/s3/telemetry/otel/DeltaHistogramTest.java b/core/src/test/java/kafka/log/stream/s3/telemetry/otel/DeltaHistogramTest.java deleted file mode 100644 index c9b1d72649..0000000000 --- a/core/src/test/java/kafka/log/stream/s3/telemetry/otel/DeltaHistogramTest.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package kafka.log.stream.s3.telemetry.otel; - -import com.yammer.metrics.core.Histogram; -import com.yammer.metrics.core.MetricsRegistry; - -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -public class DeltaHistogramTest { - - @Test - public void testDeltaMean() { - MetricsRegistry registry = new MetricsRegistry(); - Histogram histogram = registry.newHistogram(getClass(), "test-hist"); - DeltaHistogram deltaHistogram = new DeltaHistogram(histogram); - for (int i = 0; i < 10; i++) { - histogram.update(i); - } - Assertions.assertEquals(4.5, deltaHistogram.getDeltaMean()); - for (int i = 100; i < 200; i++) { - histogram.update(i); - } - Assertions.assertEquals(149.5, deltaHistogram.getDeltaMean(), 0.0001); - Assertions.assertEquals(136.31, histogram.mean(), 0.01); - } -} diff --git a/core/src/test/java/kafka/log/streamaspect/ElasticLogFileRecordsTest.java b/core/src/test/java/kafka/log/streamaspect/ElasticLogFileRecordsTest.java new file mode 100644 index 0000000000..5ea93a3dbd --- /dev/null +++ b/core/src/test/java/kafka/log/streamaspect/ElasticLogFileRecordsTest.java @@ -0,0 +1,332 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package kafka.log.streamaspect; + +import org.apache.kafka.common.compress.NoCompression; +import org.apache.kafka.common.record.MemoryRecords; +import org.apache.kafka.common.record.MemoryRecordsBuilder; +import org.apache.kafka.common.record.Record; +import org.apache.kafka.common.record.RecordBatch; +import org.apache.kafka.common.record.SimpleRecord; +import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.common.utils.ByteBufferOutputStream; +import org.apache.kafka.common.utils.Time; + +import com.automq.stream.api.Stream; +import com.automq.stream.s3.context.FetchContext; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ExecutionException; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; + +@Tag("S3Unit") +@ExtendWith(MockitoExtension.class) +class ElasticLogFileRecordsTest { + + private ElasticStreamSlice streamSlice; + private Stream stream; + + private ElasticLogFileRecords elasticLogFileRecords; + + private final Random random = new Random(); + + @BeforeEach + void setUp() { + stream = spy(new MemoryClient.StreamImpl(1)); + streamSlice = spy(new DefaultElasticStreamSlice(stream, SliceRange.of(0, Offsets.NOOP_OFFSET))); + elasticLogFileRecords = new ElasticLogFileRecords(streamSlice, 0, 0); + } + + /** + * Test reading a single, complete batch of records. + */ + @Test + void testReadSingleBatch() throws ExecutionException, InterruptedException, IOException { + // Arrange + long startOffset = 0; + int recordCount = 100; + long maxOffset = startOffset + recordCount; + final int maxReadBytes = 4096; + + Map expectedRecords = prepareRecords(startOffset, recordCount); + MemoryRecords memoryRecords = createMemoryRecords(expectedRecords); + elasticLogFileRecords.append(memoryRecords, maxOffset); + + // Act + ReadHint.setReadAll(false); + org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, maxReadBytes).get(); + + // Assert + assertNotNull(readRecords); + assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor); + assertRecords(expectedRecords, readRecords); + verify(streamSlice).fetch(any(FetchContext.class), eq(0L), eq(maxOffset), eq(maxReadBytes)); + } + + /** + * Test reading data that spans across multiple record batches. + */ + @Test + void testReadAcrossMultipleBatches() throws ExecutionException, InterruptedException, IOException { + // Arrange + long startOffset = 0; + int recordsPerBatch = 50; + int batchCount = 3; + long totalRecords = (long) recordsPerBatch * batchCount; + long maxOffset = startOffset + totalRecords; + final int maxReadBytes = Integer.MAX_VALUE; + + Map allExpectedRecords = new HashMap<>(); + long currentStartOffset = startOffset; + List nextOffsetList = new ArrayList<>(); + for (int i = 0; i < batchCount; i++) { + Map batchRecords = prepareRecords(currentStartOffset, recordsPerBatch); + allExpectedRecords.putAll(batchRecords); + MemoryRecords memoryRecords = createMemoryRecords(batchRecords); + elasticLogFileRecords.append(memoryRecords, currentStartOffset + recordsPerBatch); + + nextOffsetList.add(currentStartOffset); + currentStartOffset += recordsPerBatch; + } + + // Act + ReadHint.setReadAll(false); + org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, maxReadBytes).get(); + + // Assert + assertNotNull(readRecords); + assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor); + assertRecords(allExpectedRecords, readRecords); + nextOffsetList.forEach(nextOffset -> { + verify(streamSlice).fetch(any(FetchContext.class), eq(nextOffset), eq(maxOffset), anyInt()); + }); + } + + /** + * Test reading a batch that has been compacted, resulting in offset gaps. + * The reader should correctly iterate over the existing records and skip the gaps. + */ + @Test + void testReadCompactedBatchWithGaps() throws ExecutionException, InterruptedException, IOException { + // Arrange + long batchStartOffset = 0; + long lastOffsetInBatch = batchStartOffset + 5; // Batch spans from 0 to 5 + final int maxReadBytes = 4096; + + // Create a MemoryRecords buffer with offset gaps to simulate compaction. + // We will only include records for offsets 0, 2, 4. + Map expectedRecords = new HashMap<>(); + expectedRecords.put(batchStartOffset, createSimpleRecord("key" + batchStartOffset, "value" + batchStartOffset)); + // Skip 1 + expectedRecords.put(batchStartOffset + 2, createSimpleRecord("key" + (batchStartOffset + 2), "value" + (batchStartOffset + 2))); + // Skip 3 + expectedRecords.put(batchStartOffset + 4, createSimpleRecord("key" + (batchStartOffset + 4), "value" + (batchStartOffset + 4))); + // Skip 5 + MemoryRecords memoryRecords = createMemoryRecords(expectedRecords); + elasticLogFileRecords.append(memoryRecords, lastOffsetInBatch + 1); + + // Act: Read the entire range that contains the gappy batch. + ReadHint.setReadAll(false); + org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(batchStartOffset, lastOffsetInBatch + 1, maxReadBytes).get(); + + // Assert: Verify that only the existing records are returned. + assertNotNull(readRecords); + assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor); + assertRecords(expectedRecords, readRecords); + verify(streamSlice).fetch(any(FetchContext.class), eq(batchStartOffset), eq(lastOffsetInBatch + 1), eq(maxReadBytes)); + } + + /** + * Test reading across multiple batches where at least one batch is compacted and has offset gaps. + */ + @Test + void testReadAcrossMultipleBatchesWithGaps() throws ExecutionException, InterruptedException, IOException { + // Arrange + long startOffset = 0; + int recordsPerBatch = 50; + int batchCount = 3; + long totalRecords = (long) recordsPerBatch * batchCount; + long maxOffset = startOffset + totalRecords; + final int maxReadBytes = Integer.MAX_VALUE; + + Map allExpectedRecords = new HashMap<>(); + long currentStartOffset = startOffset; + + List nextOffsetList = new ArrayList<>(); + for (int i = 0; i < batchCount; i++) { + int skippedRecords = random.nextInt(recordsPerBatch); + Map batchRecords = prepareRecords(currentStartOffset, recordsPerBatch - skippedRecords); + allExpectedRecords.putAll(batchRecords); + MemoryRecords memoryRecords = createMemoryRecords(batchRecords); + elasticLogFileRecords.append(memoryRecords, currentStartOffset + recordsPerBatch); + + nextOffsetList.add(currentStartOffset); + currentStartOffset += recordsPerBatch; + } + + // Act + ReadHint.setReadAll(false); + org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, maxReadBytes).get(); + + // Assert + assertNotNull(readRecords); + assertTrue(readRecords instanceof ElasticLogFileRecords.BatchIteratorRecordsAdaptor); + assertRecords(allExpectedRecords, readRecords); + nextOffsetList.forEach(nextOffset -> { + verify(streamSlice).fetch(any(FetchContext.class), eq(nextOffset), eq(maxOffset), anyInt()); + }); + } + + + // Helper for preparing records to avoid code duplication + private Map prepareRecords(long startOffset, int count) { + final int fetchBatchSize = ElasticLogFileRecords.StreamSegmentInputStream.FETCH_BATCH_SIZE; + Map records = new HashMap<>(); + + // Calculate approximate size per record to ensure total size > FETCH_BATCH_SIZE + int estimatedRecordOverhead = 50; // Approximate overhead per record (headers, etc.) + int targetValueSize = Math.max(100, (fetchBatchSize / count) + estimatedRecordOverhead); + + // Create a large value string to ensure we exceed FETCH_BATCH_SIZE + StringBuilder largeValue = new StringBuilder(); + for (int j = 0; j < targetValueSize; j++) { + largeValue.append('a'); + } + String valueTemplate = largeValue.toString(); + + for (int i = 0; i < count; i++) { + long currentOffset = startOffset + i; + String value = valueTemplate + "_" + currentOffset; // Make each value unique + records.put(currentOffset, createSimpleRecord("key" + currentOffset, value)); + } + return records; + } + + // Helper for asserting records to avoid code duplication + private void assertRecords(Map expectedRecords, org.apache.kafka.common.record.Records actualRecords) { + List records = new ArrayList<>(); + actualRecords.records().forEach(records::add); + + assertEquals(expectedRecords.size(), records.size()); + + for (Record record : records) { + SimpleRecord expectedRecord = expectedRecords.get(record.offset()); + assertNotNull(expectedRecord, "Unexpected record with offset " + record.offset()); + assertEquals(new String(expectedRecord.key().array()), new String(readBytes(record.key()))); + assertEquals(new String(expectedRecord.value().array()), new String(readBytes(record.value()))); + } + } + + /** + * Test reading an empty range (startOffset >= maxOffset). + * Expects an empty Records object and no interaction with the underlying streamSlice. + */ + @Test + void testReadEmptyRange() throws ExecutionException, InterruptedException, IOException { + // Arrange + long startOffset = 100; + long maxOffset = 100; // startOffset >= maxOffset + + // Act + ReadHint.setReadAll(false); + org.apache.kafka.common.record.Records readRecords = elasticLogFileRecords.read(startOffset, maxOffset, 4096).get(); + + // Assert + assertNotNull(readRecords); + assertEquals(0, readRecords.sizeInBytes()); + // Verify that streamSlice.fetch was NOT called + verify(streamSlice, never()).fetch(any(FetchContext.class), anyLong(), anyLong(), anyInt()); + } + + // Helper methods + private SimpleRecord createSimpleRecord(String key, String value) { + return new SimpleRecord(System.currentTimeMillis(), key.getBytes(), value.getBytes()); + } + + private MemoryRecords createMemoryRecords(Map records) { + ByteBuffer buffer = ByteBuffer.allocate(128 * 1024); // Increased buffer size to accommodate larger records + long baseOffset = records.keySet().stream().min(Long::compare).get(); + ByteBufferOutputStream stream = new ByteBufferOutputStream(buffer); + try (MemoryRecordsBuilder builder = new MemoryRecordsBuilder(stream, + RecordBatch.CURRENT_MAGIC_VALUE, + NoCompression.NONE, + TimestampType.CREATE_TIME, + baseOffset, + Time.SYSTEM.milliseconds(), + RecordBatch.NO_PRODUCER_ID, + RecordBatch.NO_PRODUCER_EPOCH, + RecordBatch.NO_SEQUENCE, + false, + false, + RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.limit(), 0L)) { + + records.keySet().stream().sorted().forEach( + offset -> builder.appendWithOffset(offset, records.get(offset)) + ); + return builder.build(); + } + } + + private byte[] readBytes(ByteBuffer buffer) { + if (buffer == null) return new byte[0]; + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + return bytes; + } + + static class ReadHint { + private static final ThreadLocal READ_ALL = ThreadLocal.withInitial(() -> false); + private static final ThreadLocal FAST_READ = ThreadLocal.withInitial(() -> false); + + public static boolean isReadAll() { + return READ_ALL.get(); + } + + public static void setReadAll(boolean readAll) { + READ_ALL.set(readAll); + } + + public static boolean isFastRead() { + return FAST_READ.get(); + } + } +} diff --git a/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentManagerTest.java b/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentManagerTest.java index 9fbbbb272f..e123bb5739 100644 --- a/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentManagerTest.java +++ b/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentManagerTest.java @@ -1,18 +1,27 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * Use of this software is governed by the Business Source License - * included in the file BSL.md + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.HashSet; import java.util.Set; @@ -31,6 +40,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +@Timeout(60) @Tag("S3Unit") public class ElasticLogSegmentManagerTest { @Test diff --git a/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentTest.java b/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentTest.java index 278537352c..e5a2819334 100644 --- a/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentTest.java +++ b/core/src/test/java/kafka/log/streamaspect/ElasticLogSegmentTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -331,7 +339,7 @@ ElasticLogSegment createOrLoadSegment(long offset, int indexIntervalBytes, Time }); ElasticStreamSliceManager manager = segmentStreams.computeIfAbsent(offset, k -> { try { - return new ElasticStreamSliceManager(new ElasticLogStreamManager(new HashMap<>(), new MemoryClient.StreamClientImpl(), 1, 0, new HashMap<>())); + return new ElasticStreamSliceManager(new ElasticLogStreamManager(new HashMap<>(), new MemoryClient.StreamClientImpl(), 1, 0, new HashMap<>(), false)); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/core/src/test/java/kafka/log/streamaspect/ElasticTimeIndexTest.java b/core/src/test/java/kafka/log/streamaspect/ElasticTimeIndexTest.java index 63f5778f88..6e04579c9d 100644 --- a/core/src/test/java/kafka/log/streamaspect/ElasticTimeIndexTest.java +++ b/core/src/test/java/kafka/log/streamaspect/ElasticTimeIndexTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -20,6 +28,7 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.io.IOException; import java.util.List; @@ -27,6 +36,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; // TODO: replace S3Unit to AutoMQ +@Timeout(60) @Tag("S3Unit") public class ElasticTimeIndexTest { int maxEntries = 30; diff --git a/core/src/test/java/kafka/log/streamaspect/ElasticTransactionIndexTest.java b/core/src/test/java/kafka/log/streamaspect/ElasticTransactionIndexTest.java index 19fda7e37c..08cb34563c 100644 --- a/core/src/test/java/kafka/log/streamaspect/ElasticTransactionIndexTest.java +++ b/core/src/test/java/kafka/log/streamaspect/ElasticTransactionIndexTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -20,6 +28,7 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.io.File; import java.io.IOException; @@ -28,6 +37,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; +@Timeout(60) @Tag("S3Unit") public class ElasticTransactionIndexTest { @Test diff --git a/core/src/test/java/kafka/log/streamaspect/IStreamSliceSupplier.java b/core/src/test/java/kafka/log/streamaspect/IStreamSliceSupplier.java index 3fa17ff8c2..855d506e16 100644 --- a/core/src/test/java/kafka/log/streamaspect/IStreamSliceSupplier.java +++ b/core/src/test/java/kafka/log/streamaspect/IStreamSliceSupplier.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/test/java/kafka/log/streamaspect/PartitionStatusTrackerTest.java b/core/src/test/java/kafka/log/streamaspect/PartitionStatusTrackerTest.java index ba03111eaf..f245dff569 100644 --- a/core/src/test/java/kafka/log/streamaspect/PartitionStatusTrackerTest.java +++ b/core/src/test/java/kafka/log/streamaspect/PartitionStatusTrackerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; @@ -16,9 +24,11 @@ import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import static org.junit.jupiter.api.Assertions.assertEquals; +@Timeout(60) @Tag("S3Unit") public class PartitionStatusTrackerTest { diff --git a/core/src/test/java/kafka/log/streamaspect/cache/FileCacheTest.java b/core/src/test/java/kafka/log/streamaspect/cache/FileCacheTest.java index a7dd04c598..b1ff2e43ee 100644 --- a/core/src/test/java/kafka/log/streamaspect/cache/FileCacheTest.java +++ b/core/src/test/java/kafka/log/streamaspect/cache/FileCacheTest.java @@ -20,6 +20,7 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.io.IOException; import java.util.Arrays; @@ -32,6 +33,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; +@Timeout(60) @Tag("S3Unit") public class FileCacheTest { diff --git a/core/src/test/resources/proto/deep_nested_structure.proto b/core/src/test/resources/proto/deep_nested_structure.proto new file mode 100644 index 0000000000..6d8dd957d2 --- /dev/null +++ b/core/src/test/resources/proto/deep_nested_structure.proto @@ -0,0 +1,20 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "CompanyProtos"; + +message Company { + Department dept = 1; +} + +message Department { + Team team = 1; +} + +message Team { + repeated Member members = 1; +} + +message Member { + string name = 1; +} diff --git a/core/src/test/resources/proto/empty_collections.proto b/core/src/test/resources/proto/empty_collections.proto new file mode 100644 index 0000000000..8902913d35 --- /dev/null +++ b/core/src/test/resources/proto/empty_collections.proto @@ -0,0 +1,9 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "EmptyCasesProtos"; + +message EmptyCases { + repeated string empty_list = 1; + map empty_map = 2; +} \ No newline at end of file diff --git a/core/src/test/resources/proto/enum_type_record.proto b/core/src/test/resources/proto/enum_type_record.proto new file mode 100644 index 0000000000..1fb2c91bd3 --- /dev/null +++ b/core/src/test/resources/proto/enum_type_record.proto @@ -0,0 +1,23 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "SearchRequestProtos"; + + +enum Corpus { + CORPUS_UNSPECIFIED = 0; + CORPUS_UNIVERSAL = 1; + CORPUS_WEB = 2; + CORPUS_IMAGES = 3; + CORPUS_LOCAL = 4; + CORPUS_NEWS = 5; + CORPUS_PRODUCTS = 6; + CORPUS_VIDEO = 7; +} + +message SearchRequest { + string query = 1; + int32 page_number = 2; + int32 results_per_page = 3; + Corpus corpus = 4; +} diff --git a/core/src/test/resources/proto/list_and_nested_type.proto b/core/src/test/resources/proto/list_and_nested_type.proto new file mode 100644 index 0000000000..f026a39210 --- /dev/null +++ b/core/src/test/resources/proto/list_and_nested_type.proto @@ -0,0 +1,14 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "SearchResponseProto"; + +message SearchResponse { + repeated Result results = 1; +} + +message Result { + string url = 1; + string title = 2; + repeated string snippets = 3; +} diff --git a/core/src/test/resources/proto/map_type_record.proto b/core/src/test/resources/proto/map_type_record.proto new file mode 100644 index 0000000000..b7a2a32e6a --- /dev/null +++ b/core/src/test/resources/proto/map_type_record.proto @@ -0,0 +1,8 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "ProductProtos"; + +message Product { + map attributes = 1; +} \ No newline at end of file diff --git a/core/src/test/resources/proto/one_of_fields.proto b/core/src/test/resources/proto/one_of_fields.proto new file mode 100644 index 0000000000..802ba31a97 --- /dev/null +++ b/core/src/test/resources/proto/one_of_fields.proto @@ -0,0 +1,11 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "PaymentProtos"; + +message Payment { + oneof method { + string credit_card = 1; + string paypal = 2; + } +} \ No newline at end of file diff --git a/core/src/test/resources/proto/optional_fields_with_defaults.proto b/core/src/test/resources/proto/optional_fields_with_defaults.proto new file mode 100644 index 0000000000..2a27f50b2e --- /dev/null +++ b/core/src/test/resources/proto/optional_fields_with_defaults.proto @@ -0,0 +1,10 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "UserProto"; + +message User { + string name = 1; + optional string email = 2; // proto3的显式可选字段 + int32 age = 3; // proto3默认不设置时为0 +} \ No newline at end of file diff --git a/core/src/test/resources/proto/order.proto b/core/src/test/resources/proto/order.proto new file mode 100644 index 0000000000..dbd921d226 --- /dev/null +++ b/core/src/test/resources/proto/order.proto @@ -0,0 +1,37 @@ +syntax = "proto3"; + +package test; + +option java_package = "protobuf"; +option java_outer_classname = "OrderProtos"; + +// Order message represents a complete order +message Order { + // Address represents a physical address + message Address { + string street = 1; + string city = 2; + } + + // Customer represents the person placing the order + message Customer { + string name = 1; + Address shipping_address = 2; + Address billing_address = 3; + } + + // Price represents monetary value with currency + message Price { + double amount = 1; + string currency = 2; + } + + // Product represents an item that can be ordered + message Product { + string id = 1; + Price price = 2; + } + + Customer customer = 1; + repeated Product items = 2; +} \ No newline at end of file diff --git a/core/src/test/resources/proto/person.proto b/core/src/test/resources/proto/person.proto new file mode 100644 index 0000000000..fec898c7c4 --- /dev/null +++ b/core/src/test/resources/proto/person.proto @@ -0,0 +1,28 @@ +syntax = "proto3"; + +package kafka.automq.table.process.proto; + +import "google/protobuf/timestamp.proto"; + +option java_package = "kafka.automq.table.process.proto"; +option java_outer_classname = "PersonProto"; + +message Address { + string street = 1; + string city = 2; +} + +message Person { + int64 id = 1; + string name = 2; + bool is_active = 3; + bytes extra_data = 4; + + Address address = 5; + + repeated string roles = 6; + + map attributes = 7; + + google.protobuf.Timestamp last_updated = 8; +} diff --git a/core/src/test/resources/proto/primitive_types_record.proto b/core/src/test/resources/proto/primitive_types_record.proto new file mode 100644 index 0000000000..fec9d681a7 --- /dev/null +++ b/core/src/test/resources/proto/primitive_types_record.proto @@ -0,0 +1,22 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "TestMessageProtos"; + +message TestMessage { + string test_string = 1 [json_name = "test_str"]; + bool test_bool = 2; + bytes test_bytes = 3; + double test_double = 4; + float test_float = 5; + fixed32 test_fixed32 = 6; + fixed64 test_fixed64 = 7; + int32 test_int32 = 8; + int64 test_int64 = 9; + sfixed32 test_sfixed32 = 10; + sfixed64 test_sfixed64 = 11; + sint32 test_sint32 = 12; + sint64 test_sint64 = 13; + uint32 test_uint32 = 14; + uint64 test_uint64 = 15; +} diff --git a/core/src/test/resources/proto/product.proto b/core/src/test/resources/proto/product.proto new file mode 100644 index 0000000000..8f70b83b8d --- /dev/null +++ b/core/src/test/resources/proto/product.proto @@ -0,0 +1,27 @@ +syntax = "proto3"; + +package product; +option java_package = "protobuf.product"; +option java_outer_classname = "ProductProtos"; + +message Product { + string id = 1; + string name = 2; + string description = 3; + + enum Category { + CATEGORY_UNSPECIFIED = 0; + ELECTRONICS = 1; + CLOTHING = 2; + BOOKS = 3; + FOOD = 4; + } + + message Price { + double value = 1; + string currency_code = 2; + } + + Price price = 4; + Category category = 5; +} \ No newline at end of file diff --git a/core/src/test/resources/proto/simple_message.proto b/core/src/test/resources/proto/simple_message.proto new file mode 100644 index 0000000000..1d188cc882 --- /dev/null +++ b/core/src/test/resources/proto/simple_message.proto @@ -0,0 +1,9 @@ +syntax = "proto3"; + +package test; + +message SimpleMessage { + string name = 1; + int32 id = 2; + bool active = 3; +} diff --git a/core/src/test/resources/proto/simple_well_known_types.proto b/core/src/test/resources/proto/simple_well_known_types.proto new file mode 100644 index 0000000000..acfeb96599 --- /dev/null +++ b/core/src/test/resources/proto/simple_well_known_types.proto @@ -0,0 +1,35 @@ +syntax = "proto3"; + +package protobuf.types; + +import "google/protobuf/timestamp.proto"; +import "google/protobuf/duration.proto"; +import "google/protobuf/any.proto"; +import "google/protobuf/struct.proto"; +import "google/protobuf/wrappers.proto"; +import "google/protobuf/empty.proto"; +import "google/protobuf/field_mask.proto"; + +message ComplexMessage { + // Standard fields + string id = 1; + int32 count = 2; + + // Well-known types + google.protobuf.Timestamp created_at = 3; + google.protobuf.Duration elapsed_time = 4; + google.protobuf.Any details = 5; + google.protobuf.Struct attributes = 6; + + // Wrappers + google.protobuf.StringValue optional_name = 7; + google.protobuf.BoolValue is_active = 8; + google.protobuf.Int64Value big_count = 9; + google.protobuf.DoubleValue score = 10; + + // Additional types + google.protobuf.FieldMask update_mask = 11; + + // Empty can be used to indicate a field with no values + google.protobuf.Empty nothing = 12; +} \ No newline at end of file diff --git a/core/src/test/resources/proto/test_record.proto b/core/src/test/resources/proto/test_record.proto new file mode 100644 index 0000000000..14ed5b8541 --- /dev/null +++ b/core/src/test/resources/proto/test_record.proto @@ -0,0 +1,10 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "TestRecordProto"; + +message TestRecord { + string name = 1; + int32 age = 2; +} + diff --git a/core/src/test/resources/proto/test_record_v2.proto b/core/src/test/resources/proto/test_record_v2.proto new file mode 100644 index 0000000000..465193b6d7 --- /dev/null +++ b/core/src/test/resources/proto/test_record_v2.proto @@ -0,0 +1,11 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "TestRecordProtoV2"; + +message TestRecordV2 { + string name = 1; + int32 age = 2; + string email = 3; +} + diff --git a/core/src/test/resources/proto/timestamp_type.proto b/core/src/test/resources/proto/timestamp_type.proto new file mode 100644 index 0000000000..4c5dca924c --- /dev/null +++ b/core/src/test/resources/proto/timestamp_type.proto @@ -0,0 +1,10 @@ +syntax = "proto3"; + +option java_package = "protobuf"; +option java_outer_classname = "EventProtos"; + +import "google/protobuf/timestamp.proto"; + +message Event { + google.protobuf.Timestamp event_time = 1; +} diff --git a/core/src/test/resources/proto/with_enum.proto b/core/src/test/resources/proto/with_enum.proto new file mode 100644 index 0000000000..0ec3718ff8 --- /dev/null +++ b/core/src/test/resources/proto/with_enum.proto @@ -0,0 +1,14 @@ +syntax = "proto3"; + +package test; + +enum Status { + UNKNOWN = 0; + ACTIVE = 1; + INACTIVE = 2; +} + +message MessageWithEnum { + string name = 1; + Status status = 2; +} diff --git a/core/src/test/resources/proto/with_extensions.proto b/core/src/test/resources/proto/with_extensions.proto new file mode 100644 index 0000000000..d4b970b048 --- /dev/null +++ b/core/src/test/resources/proto/with_extensions.proto @@ -0,0 +1,21 @@ +syntax = "proto2"; + +package test; + +// Base message that allows extensions +message ExtendTestMessage { + required string id = 1; + optional string name = 2; + + // Define extension range + extensions 100 to 199; +} + +// Another message that extends the base message +message ExtendingMessage { + extend ExtendTestMessage { + optional string test_field = 100; + } + + optional string extra_field = 1; +} diff --git a/core/src/test/resources/proto/with_import.proto b/core/src/test/resources/proto/with_import.proto new file mode 100644 index 0000000000..9694b1c8e9 --- /dev/null +++ b/core/src/test/resources/proto/with_import.proto @@ -0,0 +1,10 @@ +syntax = "proto3"; + +package test; + +import "simple_message.proto"; + +message ComplexMessage { + SimpleMessage simple_message = 1; + repeated string tags = 2; +} diff --git a/core/src/test/resources/proto/with_map.proto b/core/src/test/resources/proto/with_map.proto new file mode 100644 index 0000000000..b32edb508f --- /dev/null +++ b/core/src/test/resources/proto/with_map.proto @@ -0,0 +1,22 @@ +syntax = "proto3"; +package test; + +message MapMessage { + // Various map field examples + map string_to_string = 1; + map int_to_string = 2; + map string_to_int = 3; + + // Map with complex value + message NestedValue { + string name = 1; + int32 value = 2; + bool active = 3; + } + + map string_to_nested = 4; + + // Other standard fields + string id = 5; + repeated string tags = 6; +} diff --git a/core/src/test/resources/proto/with_nested_types.proto b/core/src/test/resources/proto/with_nested_types.proto new file mode 100644 index 0000000000..1e15fbedd0 --- /dev/null +++ b/core/src/test/resources/proto/with_nested_types.proto @@ -0,0 +1,62 @@ +syntax = "proto3"; + +package test; + +// Complex message with various nested types +message ComplexNestedMessage { + // Nested enum type + enum Status { + UNKNOWN = 0; + ACTIVE = 1; + INACTIVE = 2; + PENDING = 3; + } + + // Nested message type + message Address { + string street = 1; + string city = 2; + string country = 3; + + // Deeply nested message + message GeoLocation { + double latitude = 1; + double longitude = 2; + + message Accuracy { + float precision = 1; + string source = 2; + } + Accuracy accuracy = 3; + } + GeoLocation location = 4; + } + + // Nested message with recursive type + message TreeNode { + string value = 1; + repeated TreeNode children = 2; + } + + // Using nested types in the main message + string id = 1; + Status status = 2; + Address primary_address = 3; + repeated Address secondary_addresses = 4; + TreeNode root_node = 5; + + // Message with oneof field + message ContactInfo { + oneof contact { + string email = 1; + string phone = 2; + Address physical_address = 3; + } + bool is_verified = 4; + } + repeated ContactInfo contacts = 6; + + // Nested map fields + map labeled_addresses = 7; + map labeled_nodes = 8; +} \ No newline at end of file diff --git a/core/src/test/resources/proto/with_repeated_fields.proto b/core/src/test/resources/proto/with_repeated_fields.proto new file mode 100644 index 0000000000..de88e05d36 --- /dev/null +++ b/core/src/test/resources/proto/with_repeated_fields.proto @@ -0,0 +1,26 @@ +syntax = "proto3"; +package test; + +import "google/protobuf/timestamp.proto"; + +message RepeatedFieldsMessage { + // Simple repeated fields + repeated string tags = 1; + repeated int32 values = 2; + + // Repeated complex types + message Item { + string name = 1; + int32 quantity = 2; + double price = 3; + } + + repeated Item items = 3; + + // Repeated with Google types + repeated google.protobuf.Timestamp event_times = 4; + + // Standard fields + string id = 5; + string name = 6; +} diff --git a/core/src/test/resources/proto/with_reserved.proto b/core/src/test/resources/proto/with_reserved.proto new file mode 100644 index 0000000000..d559fb7533 --- /dev/null +++ b/core/src/test/resources/proto/with_reserved.proto @@ -0,0 +1,20 @@ +syntax = "proto3"; +package test; + +message MessageWithReserved { + // Reserved field numbers + reserved 2, 15, 9 to 11, 40 to 45; + + // Reserved field names + reserved "foo", "bar", "baz"; + + // Regular fields + string id = 1; + int32 count = 3; + string name = 4; + bool active = 12; + + // Cannot use reserved fields: + // string foo = 5; // Error: field name is reserved + // int32 value = 9; // Error: field number is reserved +} diff --git a/core/src/test/resources/proto/with_service.proto b/core/src/test/resources/proto/with_service.proto new file mode 100644 index 0000000000..392ed0edf2 --- /dev/null +++ b/core/src/test/resources/proto/with_service.proto @@ -0,0 +1,74 @@ +syntax = "proto3"; +package test; + +import "google/protobuf/empty.proto"; +import "google/protobuf/timestamp.proto"; + +// Request message +message GetUserRequest { + string user_id = 1; +} + +// Response message +message User { + string id = 1; + string name = 2; + string email = 3; + google.protobuf.Timestamp created_at = 4; + bool is_active = 5; +} + +// List request +message ListUsersRequest { + int32 limit = 1; + int32 offset = 2; + string filter = 3; +} + +// List response +message ListUsersResponse { + repeated User users = 1; + int32 total = 2; + int32 next_offset = 3; +} + +// Create request +message CreateUserRequest { + string name = 1; + string email = 2; +} + +// Update request +message UpdateUserRequest { + string user_id = 1; + string name = 2; + string email = 3; + bool is_active = 4; +} + +// Service definition +service UserService { + // Get a user by ID + rpc GetUser(GetUserRequest) returns (User); + + // List users with pagination + rpc ListUsers(ListUsersRequest) returns (ListUsersResponse); + + // Create a new user + rpc CreateUser(CreateUserRequest) returns (User); + + // Update an existing user + rpc UpdateUser(UpdateUserRequest) returns (User); + + // Delete a user + rpc DeleteUser(GetUserRequest) returns (google.protobuf.Empty); + + // Stream updated users + rpc StreamUserUpdates(google.protobuf.Empty) returns (stream User); + + // Process user uploads + rpc UploadUserData(stream CreateUserRequest) returns (ListUsersResponse); + + // Bidirectional streaming + rpc ProcessUserBatch(stream UpdateUserRequest) returns (stream User); +} diff --git a/core/src/test/resources/proto/with_timestamp.proto b/core/src/test/resources/proto/with_timestamp.proto new file mode 100644 index 0000000000..bc92c139c0 --- /dev/null +++ b/core/src/test/resources/proto/with_timestamp.proto @@ -0,0 +1,11 @@ +syntax = "proto3"; +package test; + +import "google/protobuf/timestamp.proto"; + +message EventWithTimestamp { + string name = 1; + google.protobuf.Timestamp created_at = 2; + google.protobuf.Timestamp updated_at = 3; + int32 event_id = 4; +} diff --git a/core/src/test/resources/proto/with_well_known_types.proto b/core/src/test/resources/proto/with_well_known_types.proto new file mode 100644 index 0000000000..b0fb221140 --- /dev/null +++ b/core/src/test/resources/proto/with_well_known_types.proto @@ -0,0 +1,34 @@ +syntax = "proto3"; +package test; + +import "google/protobuf/timestamp.proto"; +import "google/protobuf/duration.proto"; +import "google/protobuf/any.proto"; +import "google/protobuf/struct.proto"; +import "google/protobuf/wrappers.proto"; +import "google/protobuf/empty.proto"; +import "google/protobuf/field_mask.proto"; + +message ComplexTypeMessage { + // Standard fields + string id = 1; + int32 count = 2; + + // Well-known types + google.protobuf.Timestamp created_at = 3; + google.protobuf.Duration elapsed_time = 4; + google.protobuf.Any details = 5; + google.protobuf.Struct attributes = 6; + + // Wrappers + google.protobuf.StringValue optional_name = 7; + google.protobuf.BoolValue is_active = 8; + google.protobuf.Int64Value big_count = 9; + google.protobuf.DoubleValue score = 10; + + // Additional types + google.protobuf.FieldMask update_mask = 11; + + // Empty can be used to indicate a field with no values + google.protobuf.Empty nothing = 12; +} diff --git a/core/src/test/scala/kafka/log/streamaspect/ElasticLogCleanerTest.scala b/core/src/test/scala/kafka/log/streamaspect/ElasticLogCleanerTest.scala index a865be96a4..9b30a8db25 100644 --- a/core/src/test/scala/kafka/log/streamaspect/ElasticLogCleanerTest.scala +++ b/core/src/test/scala/kafka/log/streamaspect/ElasticLogCleanerTest.scala @@ -8,12 +8,13 @@ import org.apache.kafka.common.Uuid import org.apache.kafka.common.config.TopicConfig import org.apache.kafka.coordinator.transaction.TransactionLogConfigs import org.apache.kafka.storage.internals.log.{FetchIsolation, LogConfig, LogDirFailureChannel, LogOffsetsListener} -import org.junit.jupiter.api.{Assertions, BeforeEach, Tag, Test} +import org.junit.jupiter.api.{Assertions, BeforeEach, Tag, Test, Timeout} import java.io.File import java.util.Properties import scala.jdk.CollectionConverters.CollectionHasAsScala +@Timeout(60) @Tag("S3Unit") class ElasticLogCleanerTest extends LogCleanerTest { var client: Client = _ @@ -80,6 +81,51 @@ class ElasticLogCleanerTest extends LogCleanerTest { } } + @Test + @Timeout(value = 30) + def testCleanSegmentCauseHollowWithEmptySegment(): Unit = { + val cleaner = makeCleaner(Int.MaxValue) + val logProps = new Properties() + logProps.put(TopicConfig.SEGMENT_BYTES_CONFIG, 1024: java.lang.Integer) + + val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps)) + log.appendAsLeader(record(1, 0), leaderEpoch = 0) + log.appendAsLeader(record(2, 0), leaderEpoch = 0) + while (log.numberOfSegments < 2) { + log.appendAsLeader(record(1, log.logEndOffset.toInt), leaderEpoch = 0) + } + while (log.numberOfSegments < 3) { + log.appendAsLeader(record(3, 22), leaderEpoch = 0) + } + log.appendAsLeader(record(1, log.logEndOffset.toInt), leaderEpoch = 0) + log.appendAsLeader(record(3, log.logEndOffset.toInt), leaderEpoch = 0) + + val map = new FakeOffsetMap(Int.MaxValue) + map.put(key(2L), 1L) + map.put(key(1L), log.logEndOffset - 2) + map.put(key(3L), log.logEndOffset - 1) + + // create an empty segment in between first and last segment + cleaner.cleanSegments(log, log.logSegments.asScala.take(1).toSeq, map, 0L, new CleanerStats, new CleanedTransactionMetadata, -1) + cleaner.cleanSegments(log, log.logSegments.asScala.slice(1, 2).toSeq, map, 0L, new CleanerStats, new CleanedTransactionMetadata, -1) + + log.logSegments.asScala.slice(1, 2).foreach(s => { + Assertions.assertEquals(0, s.size()) + }) + + var offset = 0L + var total = 0 + while (offset < log.logEndOffset) { + val rst = log.read(offset, 1, FetchIsolation.LOG_END, minOneMessage = true) + Assertions.assertNotNull(rst) + rst.records.batches.forEach(b => { + total += 1 + offset = b.nextOffset() + }) + } + Assertions.assertEquals(4, total) + } + override protected def makeLog(dir: File, config: LogConfig, recoveryPoint: Long): ElasticUnifiedLog = { diff --git a/core/src/test/scala/kafka/log/streamaspect/ElasticLogLoaderTest.java b/core/src/test/scala/kafka/log/streamaspect/ElasticLogLoaderTest.java index 393012f1ee..3498a2e046 100644 --- a/core/src/test/scala/kafka/log/streamaspect/ElasticLogLoaderTest.java +++ b/core/src/test/scala/kafka/log/streamaspect/ElasticLogLoaderTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect; diff --git a/core/src/test/scala/kafka/log/streamaspect/ElasticLogTest.scala b/core/src/test/scala/kafka/log/streamaspect/ElasticLogTest.scala index 9b35925249..a545af4ff8 100644 --- a/core/src/test/scala/kafka/log/streamaspect/ElasticLogTest.scala +++ b/core/src/test/scala/kafka/log/streamaspect/ElasticLogTest.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect @@ -24,7 +32,7 @@ import org.apache.kafka.coordinator.transaction.TransactionLogConfigs import org.apache.kafka.server.util.{MockTime, Scheduler} import org.apache.kafka.storage.internals.log._ import org.junit.jupiter.api.Assertions._ -import org.junit.jupiter.api.{AfterEach, BeforeEach, Tag, Test} +import org.junit.jupiter.api.{AfterEach, BeforeEach, Tag, Test, Timeout} import java.io.{File, IOException} import java.nio.charset.StandardCharsets @@ -33,6 +41,7 @@ import java.util.regex.Pattern import scala.jdk.CollectionConverters.IterableHasAsScala // TODO: extends the LocalLogTest +@Timeout(60) @Tag("S3Unit") class ElasticLogTest { val kafkaConfig: KafkaConfig = KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, "127.0.0.1:1", port = -1)) diff --git a/core/src/test/scala/kafka/log/streamaspect/ElasticProducerStateManagerTest.scala b/core/src/test/scala/kafka/log/streamaspect/ElasticProducerStateManagerTest.scala index f9e9f1d558..a4b20941b2 100644 --- a/core/src/test/scala/kafka/log/streamaspect/ElasticProducerStateManagerTest.scala +++ b/core/src/test/scala/kafka/log/streamaspect/ElasticProducerStateManagerTest.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect @@ -20,11 +28,10 @@ import org.apache.kafka.common.utils.{MockTime, Utils} import org.apache.kafka.coordinator.transaction.TransactionLogConfigs import org.apache.kafka.storage.internals.log._ import org.junit.jupiter.api.Assertions._ -import org.junit.jupiter.api.{AfterEach, BeforeEach, Tag, Test} +import org.junit.jupiter.api.{AfterEach, BeforeEach, Tag, Test, Timeout} import org.junit.jupiter.params.ParameterizedTest import org.junit.jupiter.params.provider.ValueSource import org.mockito.Mockito.{mock, when} - import kafka.log.streamaspect.ElasticProducerStateManager.AWAIT_SEQ_ZERO_TIMEOUT import java.io.File @@ -36,6 +43,7 @@ import java.util.{Collections, Optional, OptionalLong} import scala.compat.java8.OptionConverters.RichOptionalGeneric import scala.jdk.CollectionConverters._ +@Timeout(60) @Tag("S3Unit") class ElasticProducerStateManagerTest { diff --git a/core/src/test/scala/kafka/log/streamaspect/ElasticUnifiedLogTest.scala b/core/src/test/scala/kafka/log/streamaspect/ElasticUnifiedLogTest.scala index 1da0985897..f7833e183d 100644 --- a/core/src/test/scala/kafka/log/streamaspect/ElasticUnifiedLogTest.scala +++ b/core/src/test/scala/kafka/log/streamaspect/ElasticUnifiedLogTest.scala @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.log.streamaspect @@ -25,7 +33,7 @@ import org.apache.kafka.server.util.Scheduler import org.apache.kafka.storage.internals.log._ import org.apache.kafka.storage.internals.utils.Throttler import org.junit.jupiter.api.Assertions._ -import org.junit.jupiter.api.{BeforeEach, Tag, Test} +import org.junit.jupiter.api.{BeforeEach, Tag, Test, Timeout} import org.junit.jupiter.params.ParameterizedTest import org.junit.jupiter.params.provider.ValueSource @@ -33,6 +41,7 @@ import java.io.File import java.util.Optional import scala.jdk.CollectionConverters.IterableHasAsScala +@Timeout(60) @Tag("S3Unit") class ElasticUnifiedLogTest extends UnifiedLogTest { var client: Client = _ diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala index 35e697827e..8ea0eb44ef 100644 --- a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala +++ b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala @@ -30,7 +30,7 @@ import org.apache.kafka.common.utils.Utils import org.apache.kafka.coordinator.transaction.TransactionLogConfigs import org.apache.kafka.server.metrics.{KafkaMetricsGroup, KafkaYammerMetrics} import org.apache.kafka.server.util.MockTime -import org.apache.kafka.storage.internals.log.{AbortedTxn, AppendOrigin, CleanerConfig, LogAppendInfo, LogConfig, LogDirFailureChannel, LogFileUtils, LogSegment, LogSegments, LogStartOffsetIncrementReason, OffsetMap, ProducerStateManager, ProducerStateManagerConfig} +import org.apache.kafka.storage.internals.log._ import org.apache.kafka.storage.internals.utils.Throttler import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.{AfterEach, Test} @@ -1437,8 +1437,20 @@ class LogCleanerTest extends Logging { //create 3 segments for (i <- 0 until 3) { log.appendAsLeader(TestUtils.singletonRecords(value = v, key = k), leaderEpoch = 0) - //0 to Int.MaxValue is Int.MaxValue+1 message, -1 will be the last message of i-th segment - val records = messageWithOffset(k, v, (i + 1L) * (Int.MaxValue + 1L) -1 ) + + // AutoMQ inject start + val records = if (log.isInstanceOf[ElasticUnifiedLog]) { + // Create a sparse segment by appending a record with a large offset. + // A segment can contain up to Int.MaxValue messages (see https://github.com/AutoMQ/automq/issues/2717). + // The offset `(i + 1L) * Int.MaxValue - 1` ensures that this is the last message of the i-th segment, + // creating a large gap to the next segment's base offset. This helps test segment grouping with sparse offsets. + messageWithOffset(k, v, (i + 1L) * Int.MaxValue - 1) + } else { + //0 to Int.MaxValue is Int.MaxValue+1 message, -1 will be the last message of i-th segment + messageWithOffset(k, v, (i + 1L) * (Int.MaxValue + 1L) -1 ) + } + + // AutoMQ inject end log.appendAsFollower(records) assertEquals(i + 1, log.numberOfSegments) } @@ -1469,6 +1481,13 @@ class LogCleanerTest extends Logging { //trigger a clean and 2 empty segments should cleaned to 1 cleaner.clean(LogToClean(log.topicPartition, log, 0, firstUncleanableOffset)) assertEquals(totalSegments - 1, log.numberOfSegments) + + // AutoMQ inject start + + // after clean, the 2nd and 3rd segment should be none empty + assertEquals(2, log.logSegments.asScala.takeRight(2).count(_.size > 0)) + + // AutoMQ inject end } /** @@ -1492,12 +1511,21 @@ class LogCleanerTest extends Logging { log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, key = "hello".getBytes), leaderEpoch = 0) // forward offset and append message to next segment at offset Int.MaxValue - val records = messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 1) + // AutoMQ inject start + val records = if (log.isInstanceOf[ElasticUnifiedLog]) { + messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 2) + } else { + messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 1) + } + // AutoMQ inject end + log.appendAsFollower(records) log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, key = "hello".getBytes), leaderEpoch = 0) // AutoMQ inject start - if (!log.isInstanceOf[ElasticUnifiedLog]) { + if (log.isInstanceOf[ElasticUnifiedLog]) { + assertEquals(Int.MaxValue - 1, log.activeSegment.readNextOffset() - 1) + } else { assertEquals(Int.MaxValue, log.activeSegment.offsetIndex.lastOffset) } // AutoMQ inject end diff --git a/core/src/test/scala/unit/kafka/server/BrokerQuotaManagerTest.java b/core/src/test/scala/unit/kafka/server/BrokerQuotaManagerTest.java index aa2bf11586..ecbeff96d8 100644 --- a/core/src/test/scala/unit/kafka/server/BrokerQuotaManagerTest.java +++ b/core/src/test/scala/unit/kafka/server/BrokerQuotaManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package unit.kafka.server; @@ -29,6 +37,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.Properties; import java.util.concurrent.TimeUnit; @@ -41,6 +50,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +@Timeout(60) @Tag("s3Unit") public class BrokerQuotaManagerTest { private final Time time = Time.SYSTEM; @@ -105,27 +115,115 @@ public void testQuota() { result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.fetch(), request, 500, time + second2millis); assertEquals(0, result); - // Test request quota + // Test slow fetch quota properties.put(QuotaConfigs.BROKER_QUOTA_FETCH_BYTES_CONFIG, 0); + properties.put(QuotaConfigs.BROKER_QUOTA_SLOW_FETCH_BYTES_CONFIG, 100); + brokerQuotaManager.updateQuotaConfigs(Option.apply(properties)); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.slowFetch(), request, 100, time); + assertEquals(0, result); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.slowFetch(), request, 100, time + 10); + assertEquals(0, result); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.slowFetch(), request, 100, time + second2millis); + assertTrue(result > 0); + + properties.put(QuotaConfigs.BROKER_QUOTA_SLOW_FETCH_BYTES_CONFIG, 1000); + brokerQuotaManager.updateQuotaConfigs(Option.apply(properties)); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.slowFetch(), request, 500, time + second2millis); + assertEquals(0, result); + + // Test request quota + properties.put(QuotaConfigs.BROKER_QUOTA_SLOW_FETCH_BYTES_CONFIG, 0); properties.put(QuotaConfigs.BROKER_QUOTA_REQUEST_RATE_CONFIG, 1); brokerQuotaManager.updateQuotaConfigs(Option.apply(properties)); - result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.request(), request, 1, time); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time); assertEquals(0, result); - result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.request(), request, 1, time + 10); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 10); assertEquals(0, result); - result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.request(), request, 1, time + second2millis); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + second2millis); assertTrue(result > 0); properties.put(QuotaConfigs.BROKER_QUOTA_REQUEST_RATE_CONFIG, 10); brokerQuotaManager.updateQuotaConfigs(Option.apply(properties)); - result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.request(), request, 0, time + second2millis); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 0, time + second2millis); assertEquals(0, result); } + @Test + public void testZeroQuota() { + long result; + long time = this.time.milliseconds(); + + // enable quota + Properties properties = new Properties(); + properties.put(QuotaConfigs.BROKER_QUOTA_ENABLED_CONFIG, true); + brokerQuotaManager.updateQuotaConfigs(Option.apply(properties)); + + brokerQuotaManager.updateQuota(QuotaType.requestRate(), 0); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time); + assertEquals(1000, result); + + brokerQuotaManager.updateQuota(QuotaType.slowFetch(), 0); + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.slowFetch(), request, 1, time); + assertEquals(1000, result); + } + + @Test + public void testUpdateQuota() { + int result; + long time = this.time.milliseconds(); + + // enable quota + Properties properties = new Properties(); + properties.put(QuotaConfigs.BROKER_QUOTA_ENABLED_CONFIG, true); + brokerQuotaManager.updateQuotaConfigs(Option.apply(properties)); + + brokerQuotaManager.updateQuota(QuotaType.requestRate(), 1); + // rate = 1 / 2000ms + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 1 / 2, time); + assertEquals(0, result); + // rate = 2 / 2010ms + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 10); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 2 / 2.01, time + 10); + assertEquals(0, result); + // rate = 3 / 2999ms > 1 + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 3 / 2.999, time + 2999); + assertEquals(1, result); + + brokerQuotaManager.updateQuota(QuotaType.requestRate(), 2); + // rate = 4 / 2999ms + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 4 / 2.999, time + 2999); + assertEquals(0, result); + // rate = 5 / 2999ms + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 5 / 2.999, time + 2999); + assertEquals(0, result); + // rate = 6 / 2999ms > 2 + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 6 / 2.999, time + 2999); + assertEquals(1, result); + + brokerQuotaManager.updateQuota(QuotaType.requestRate(), 1); + // rate = 5 / 2999ms > 1 + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999 + 2999); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 5 / 2.999, time + 2999 + 2999); + assertEquals(1000, result); + // rate = 2 / 2001ms + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999 + 2999 + 1); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 2 / 2.001, time + 2999 + 2999 + 1); + assertEquals(0, result); + // rate = 3 / 2999ms > 1 + result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.requestRate(), request, 1, time + 2999 + 2999 + 2999); + assertQuotaMetricValue(QuotaType.requestRate(), (double) 3 / 2.999, time + 2999 + 2999 + 2999); + assertEquals(1, result); + } + @Test public void testThrottle() { AtomicInteger throttleCounter = new AtomicInteger(0); - brokerQuotaManager.throttle(QuotaType.request(), new ThrottleCallback() { + brokerQuotaManager.throttle(QuotaType.requestRate(), new ThrottleCallback() { @Override public void startThrottling() { throttleCounter.incrementAndGet(); @@ -184,4 +282,9 @@ public void testWhiteList() { result = brokerQuotaManager.maybeRecordAndGetThrottleTimeMs(QuotaType.produce(), request, 1000, time.milliseconds()); assertEquals(0, result); } + + private void assertQuotaMetricValue(QuotaType quotaType, double expected, long timeMs) { + double value = brokerQuotaManager.getQuotaMetricValue(quotaType, timeMs).get(); + assertEquals(expected, value); + } } diff --git a/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTableTest.scala b/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTableTest.scala new file mode 100644 index 0000000000..982451329a --- /dev/null +++ b/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTableTest.scala @@ -0,0 +1,115 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package unit.kafka.server + +import kafka.automq.AutoMQConfig +import kafka.server.{ControllerConfigurationValidator, KafkaConfig} +import kafka.utils.TestUtils +import org.apache.kafka.common.config.ConfigResource +import org.apache.kafka.common.config.ConfigResource.Type.TOPIC +import org.apache.kafka.common.config.TopicConfig.{AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG, AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG} +import org.apache.kafka.common.errors.InvalidConfigurationException +import org.apache.kafka.server.record.{TableTopicConvertType, TableTopicTransformType} +import org.junit.jupiter.api.Assertions.{assertDoesNotThrow, assertEquals, assertThrows} +import org.junit.jupiter.api.{BeforeEach, Tag, Test, Timeout} + +import java.util +import java.util.Locale + +@Timeout(60) +@Tag("S3Unit") +class ControllerConfigurationValidatorTableTest { + + private var validator: ControllerConfigurationValidator = _ + private var validatorWithSchemaRegistry: ControllerConfigurationValidator = _ + + @BeforeEach + def setUp(): Unit = { + val config = new KafkaConfig(TestUtils.createDummyBrokerConfig()) + validator = new ControllerConfigurationValidator(config) + + val brokerConfigWithSchemaRegistry = TestUtils.createDummyBrokerConfig() + brokerConfigWithSchemaRegistry.put(AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG, "http://localhost:8081") + val kafkaConfigWithSchemaRegistry = new KafkaConfig(brokerConfigWithSchemaRegistry) + validatorWithSchemaRegistry = new ControllerConfigurationValidator(kafkaConfigWithSchemaRegistry) + } + + @Test + def testConvertTypeWithSchemaRegistryUrlNotConfigured(): Unit = { + val config = new util.TreeMap[String, String]() + config.put(AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG, TableTopicTransformType.NONE.name) + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, TableTopicConvertType.BY_SCHEMA_ID.name) + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG, TableTopicConvertType.STRING.name) + + var exception = assertThrows(classOf[InvalidConfigurationException], () => { + validator.validate(new ConfigResource(TOPIC, "foo"), config) + }) + assertEquals("Table topic convert type is set to 'by_schema_id' but schema registry URL is not configured", exception.getMessage) + + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, TableTopicConvertType.BY_LATEST_SCHEMA.name) + exception = assertThrows(classOf[InvalidConfigurationException], () => { + validator.validate(new ConfigResource(TOPIC, "foo"), config) + }) + assertEquals("Table topic convert type is set to 'by_latest_schema' but schema registry URL is not configured", exception.getMessage) + } + + @Test + def testConvertTypeWithSchemaRegistryUrlConfigured(): Unit = { + val config = new util.TreeMap[String, String]() + config.put(AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG, TableTopicTransformType.NONE.name) + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG, TableTopicConvertType.STRING.name) + + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, TableTopicConvertType.BY_SCHEMA_ID.name) + assertDoesNotThrow(new org.junit.jupiter.api.function.Executable { def execute(): Unit = validatorWithSchemaRegistry.validate(new ConfigResource(TOPIC, "foo"), config) }) + + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, TableTopicConvertType.BY_LATEST_SCHEMA.name) + assertDoesNotThrow(new org.junit.jupiter.api.function.Executable { def execute(): Unit = validatorWithSchemaRegistry.validate(new ConfigResource(TOPIC, "foo"), config) }) + } + + @Test + def testRawConvertTypeWithDebeziumUnwrapTransform(): Unit = { + val config = new util.TreeMap[String, String]() + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, TableTopicConvertType.RAW.name) + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG, TableTopicConvertType.STRING.name) + config.put(AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG, TableTopicTransformType.FLATTEN_DEBEZIUM.name().toLowerCase(Locale.ROOT)) + + val exception = assertThrows(classOf[InvalidConfigurationException], () => { + validator.validate(new ConfigResource(TOPIC, "foo"), config) + }) + assertEquals("raw convert type cannot be used with 'flatten_debezium' transform type", exception.getMessage) + } + + @Test + def testValidRawConvertType(): Unit = { + val config = new util.TreeMap[String, String]() + config.put(AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG, TableTopicTransformType.NONE.name) + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, TableTopicConvertType.RAW.name) + config.put(AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG, TableTopicConvertType.STRING.name) + + assertDoesNotThrow(new org.junit.jupiter.api.function.Executable { def execute(): Unit = validator.validate(new ConfigResource(TOPIC, "foo"), config) }) + } + + @Test + def testEmptyTableTopicConfigShouldBeValid(): Unit = { + val config = new util.TreeMap[String, String]() + + assertDoesNotThrow(new org.junit.jupiter.api.function.Executable { def execute(): Unit = validator.validate(new ConfigResource(TOPIC, "foo"), config) }) + } +} diff --git a/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTest.scala b/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTest.scala index 00bb93811b..29ea7ec58a 100644 --- a/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTest.scala +++ b/core/src/test/scala/unit/kafka/server/ControllerConfigurationValidatorTest.scala @@ -17,12 +17,14 @@ package kafka.server +import kafka.automq.AutoMQConfig import kafka.utils.TestUtils import org.apache.kafka.common.config.ConfigResource import org.apache.kafka.common.config.ConfigResource.Type.{BROKER, BROKER_LOGGER, CLIENT_METRICS, TOPIC} -import org.apache.kafka.common.config.TopicConfig.{SEGMENT_BYTES_CONFIG, SEGMENT_JITTER_MS_CONFIG, SEGMENT_MS_CONFIG} +import org.apache.kafka.common.config.TopicConfig.{SEGMENT_BYTES_CONFIG, SEGMENT_JITTER_MS_CONFIG, SEGMENT_MS_CONFIG, TABLE_TOPIC_SCHEMA_TYPE_CONFIG} import org.apache.kafka.common.errors.{InvalidConfigurationException, InvalidRequestException, InvalidTopicException} import org.apache.kafka.server.metrics.ClientMetricsConfigs +import org.apache.kafka.server.record.TableTopicSchemaType import org.junit.jupiter.api.Assertions.{assertEquals, assertThrows} import org.junit.jupiter.api.Test @@ -153,4 +155,26 @@ class ControllerConfigurationValidatorTest { assertThrows(classOf[InvalidConfigurationException], () => validator.validate( new ConfigResource(CLIENT_METRICS, "subscription-1"), config)). getMessage) } + + @Test + def testInvalidTableTopicSchemaConfig(): Unit = { + val config = new util.TreeMap[String, String]() + config.put(TABLE_TOPIC_SCHEMA_TYPE_CONFIG, TableTopicSchemaType.SCHEMA.name) + + // Test without schema registry URL configured + val exception = assertThrows(classOf[InvalidRequestException], () => { + validator.validate(new ConfigResource(TOPIC, "foo"), config) + }) + assertEquals("Table topic schema type is set to SCHEMA but schema registry URL is not configured", exception.getMessage) + + // Test with schema registry URL configured + val brokerConfigWithSchemaRegistry = TestUtils.createDummyBrokerConfig() + brokerConfigWithSchemaRegistry.put(AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG, "http://localhost:8081") + + val kafkaConfigWithSchemaRegistry = new KafkaConfig(brokerConfigWithSchemaRegistry) + val validatorWithSchemaRegistry = new ControllerConfigurationValidator(kafkaConfigWithSchemaRegistry) + + // No exception should be thrown when schema registry URL is configured properly + validatorWithSchemaRegistry.validate(new ConfigResource(TOPIC, "foo"), config) + } } diff --git a/core/src/test/scala/unit/kafka/server/streamaspect/ElasticControllerApisTest.scala b/core/src/test/scala/unit/kafka/server/streamaspect/ElasticControllerApisTest.scala index 10aca297c2..73024a150f 100644 --- a/core/src/test/scala/unit/kafka/server/streamaspect/ElasticControllerApisTest.scala +++ b/core/src/test/scala/unit/kafka/server/streamaspect/ElasticControllerApisTest.scala @@ -9,10 +9,11 @@ import org.apache.kafka.raft.QuorumConfig import org.apache.kafka.server.authorizer.Authorizer import org.apache.kafka.server.common.{FinalizedFeatures, MetadataVersion} import org.apache.kafka.server.config.KRaftConfigs -import org.junit.jupiter.api.Tag +import org.junit.jupiter.api.{Tag, Timeout} import java.util.Properties +@Timeout(60) @Tag("S3Unit") class ElasticControllerApisTest extends ControllerApisTest { diff --git a/core/src/test/scala/unit/kafka/server/streamaspect/ElasticKafkaApisTest.scala b/core/src/test/scala/unit/kafka/server/streamaspect/ElasticKafkaApisTest.scala index 0cdbb3efd7..3b33a6acc3 100644 --- a/core/src/test/scala/unit/kafka/server/streamaspect/ElasticKafkaApisTest.scala +++ b/core/src/test/scala/unit/kafka/server/streamaspect/ElasticKafkaApisTest.scala @@ -11,13 +11,14 @@ import org.apache.kafka.raft.QuorumConfig import org.apache.kafka.server.authorizer.Authorizer import org.apache.kafka.server.common.{FinalizedFeatures, MetadataVersion} import org.apache.kafka.server.config.KRaftConfigs -import org.junit.jupiter.api.Tag +import org.junit.jupiter.api.{Tag, Timeout} import org.mockito.Mockito.mock import java.util.Collections import scala.collection.Map import scala.jdk.CollectionConverters.SetHasAsScala +@Timeout(60) @Tag("S3Unit") class ElasticKafkaApisTest extends KafkaApisTest { override protected val replicaManager: ElasticReplicaManager = mock(classOf[ElasticReplicaManager]) diff --git a/core/src/test/scala/unit/kafka/server/streamaspect/ElasticReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/streamaspect/ElasticReplicaManagerTest.scala index 0e5a303f4a..b611f3d528 100644 --- a/core/src/test/scala/unit/kafka/server/streamaspect/ElasticReplicaManagerTest.scala +++ b/core/src/test/scala/unit/kafka/server/streamaspect/ElasticReplicaManagerTest.scala @@ -30,7 +30,7 @@ import org.apache.kafka.server.util.timer.MockTimer import org.apache.kafka.server.util.{MockScheduler, Scheduler} import org.apache.kafka.storage.internals.log._ import org.junit.jupiter.api.Assertions.assertEquals -import org.junit.jupiter.api.{BeforeEach, Disabled, Tag, Test} +import org.junit.jupiter.api.{BeforeEach, Disabled, Tag, Test, Timeout} import org.junit.jupiter.params.ParameterizedTest import org.junit.jupiter.params.provider.ValueSource import org.mockito.ArgumentMatchers @@ -46,6 +46,7 @@ import scala.collection.{Map, Seq} import scala.compat.java8.OptionConverters.RichOptionForJava8 import scala.jdk.CollectionConverters.{MapHasAsJava, PropertiesHasAsScala} +@Timeout(60) @Tag("S3Unit") class ElasticReplicaManagerTest extends ReplicaManagerTest { @@ -93,14 +94,12 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { threadNamePrefix: Option[String] = None, brokerEpochSupplier: () => Long = () => -1, addPartitionsToTxnManager: Option[AddPartitionsToTxnManager] = None, - directoryEventHandler: DirectoryEventHandler = DirectoryEventHandler.NOOP): ElasticReplicaManager = { - new ElasticReplicaManager(config, metrics, time, scheduler, logManager, remoteLogManager, quotaManagers, + directoryEventHandler: DirectoryEventHandler = DirectoryEventHandler.NOOP) = new ElasticReplicaManager(config, metrics, time, scheduler, logManager, remoteLogManager, quotaManagers, metadataCache, logDirFailureChannel, alterPartitionManager, brokerTopicStats, isShuttingDown, zkClient, delayedProducePurgatoryParam, delayedFetchPurgatoryParam, delayedDeleteRecordsPurgatoryParam, delayedElectLeaderPurgatoryParam, delayedRemoteFetchPurgatoryParam, threadNamePrefix, brokerEpochSupplier, addPartitionsToTxnManager, directoryEventHandler, MoreExecutors.newDirectExecutorService(), MoreExecutors.newDirectExecutorService()) - } override protected def setUpReplicaManagerWithMockedAddPartitionsToTxnManager(addPartitionsToTxnManager: AddPartitionsToTxnManager, transactionalTopicPartitions: List[TopicPartition], @@ -149,15 +148,11 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { if (enableRemoteStorage) { props.put("log.dirs", path1) props.put(RemoteLogManagerConfig.REMOTE_LOG_STORAGE_SYSTEM_ENABLE_PROP, enableRemoteStorage.toString) - } else { - props.put("log.dirs", path1 + "," + path2) - } + } else props.put("log.dirs", path1 + "," + path2) propsModifier.apply(props) val config = KafkaConfig.fromProps(props) val logProps = new Properties() - if (enableRemoteStorage && defaultTopicRemoteLogStorageEnable) { - logProps.put(TopicConfig.REMOTE_LOG_STORAGE_ENABLE_CONFIG, "true") - } + if (enableRemoteStorage && defaultTopicRemoteLogStorageEnable) logProps.put(TopicConfig.REMOTE_LOG_STORAGE_ENABLE_CONFIG, "true") val mockLog = setupMockLog(path1) if (setupLogDirMetaProperties) { // add meta.properties file in each dir @@ -180,7 +175,7 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { val aliveBrokers = aliveBrokerIds.map(brokerId => new Node(brokerId, s"host$brokerId", brokerId)) brokerTopicStats = new BrokerTopicStats() - val metadataCache: MetadataCache = mock(classOf[MetadataCache]) + val metadataCache = mock(classOf[MetadataCache]) when(metadataCache.topicIdInfo()).thenReturn((topicIds.asJava, topicNames.asJava)) when(metadataCache.topicNamesToIds()).thenReturn(topicIds.asJava) when(metadataCache.topicIdsToNames()).thenReturn(topicNames.asJava) @@ -226,12 +221,10 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { threadNamePrefix = Option(this.getClass.getName), addPartitionsToTxnManager = Some(addPartitionsToTxnManager), directoryEventHandler = directoryEventHandler, - remoteLogManager = if (enableRemoteStorage) { - if (remoteLogManager.isDefined) - remoteLogManager - else - Some(mockRemoteLogManager) - } else None, + remoteLogManager = if (enableRemoteStorage) if (remoteLogManager.isDefined) + remoteLogManager + else + Some(mockRemoteLogManager) else None, fastFetchExecutor = MoreExecutors.newDirectExecutorService(), slowFetchExecutor = MoreExecutors.newDirectExecutorService()) { @@ -240,66 +233,58 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { time: Time, threadNamePrefix: Option[String], quotaManager: ReplicationQuotaManager - ): ReplicaFetcherManager = { - mockReplicaFetcherManager.getOrElse { - if (buildRemoteLogAuxState) { - super.createReplicaFetcherManager( - metrics, - time, - threadNamePrefix, - quotaManager - ) - val config = this.config - val metadataCache = this.metadataCache - new ReplicaFetcherManager(config, this, metrics, time, threadNamePrefix, quotaManager, () => metadataCache.metadataVersion(), () => 1) { - override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaFetcherThread = { - val prefix = threadNamePrefix.map(tp => s"$tp:").getOrElse("") - val threadName = s"${prefix}ReplicaFetcherThread-$fetcherId-${sourceBroker.id}" - - val tp = new TopicPartition(topic, 0) - val leader = new MockLeaderEndPoint() { - override def fetch(fetchRequest: FetchRequest.Builder): Map[TopicPartition, FetchData] = { - Map(tp -> new FetchData().setErrorCode(Errors.OFFSET_MOVED_TO_TIERED_STORAGE.code)) - } - } - leader.setLeaderState(tp, PartitionState(leaderEpoch = 0)) - leader.setReplicaPartitionStateCallback(tp => PartitionState(leaderEpoch = 0)) - - val fetcher = new ReplicaFetcherThread(threadName, leader, config, failedPartitions, replicaManager, - quotaManager, "", () => config.interBrokerProtocolVersion) - - val initialFetchState = InitialFetchState( - topicId = Some(Uuid.randomUuid()), - leader = leader.brokerEndPoint(), - currentLeaderEpoch = 0, - initOffset = 0) - - fetcher.addPartitions(Map(tp -> initialFetchState)) - - fetcher + ): ReplicaFetcherManager = mockReplicaFetcherManager.getOrElse { + if (buildRemoteLogAuxState) { + super.createReplicaFetcherManager( + metrics, + time, + threadNamePrefix, + quotaManager + ) + val config = this.config + val metadataCache = this.metadataCache + new ReplicaFetcherManager(config, this, metrics, time, threadNamePrefix, quotaManager, () => metadataCache.metadataVersion(), () => 1) { + override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaFetcherThread = { + val prefix = threadNamePrefix.map(tp => s"$tp:").getOrElse("") + val threadName = s"${prefix}ReplicaFetcherThread-$fetcherId-${sourceBroker.id}" + + val tp = new TopicPartition(topic, 0) + val leader = new MockLeaderEndPoint() { + override def fetch(fetchRequest: FetchRequest.Builder): Map[TopicPartition, FetchData] = Map(tp -> new FetchData().setErrorCode(Errors.OFFSET_MOVED_TO_TIERED_STORAGE.code)) } + leader.setLeaderState(tp, PartitionState(leaderEpoch = 0)) + leader.setReplicaPartitionStateCallback(tp => PartitionState(leaderEpoch = 0)) + + val fetcher = new ReplicaFetcherThread(threadName, leader, config, failedPartitions, replicaManager, + quotaManager, "", () => config.interBrokerProtocolVersion) + + val initialFetchState = InitialFetchState( + topicId = Some(Uuid.randomUuid()), + leader = leader.brokerEndPoint(), + currentLeaderEpoch = 0, + initOffset = 0) + + fetcher.addPartitions(Map(tp -> initialFetchState)) + + fetcher } - } else { - super.createReplicaFetcherManager( - metrics, - time, - threadNamePrefix, - quotaManager - ) } - } + } else super.createReplicaFetcherManager( + metrics, + time, + threadNamePrefix, + quotaManager + ) } override def createReplicaAlterLogDirsManager( quotaManager: ReplicationQuotaManager, brokerTopicStats: BrokerTopicStats - ): ReplicaAlterLogDirsManager = { - mockReplicaAlterLogDirsManager.getOrElse { - super.createReplicaAlterLogDirsManager( - quotaManager, - brokerTopicStats - ) - } + ): ReplicaAlterLogDirsManager = mockReplicaAlterLogDirsManager.getOrElse { + super.createReplicaAlterLogDirsManager( + quotaManager, + brokerTopicStats + ) } } } @@ -375,7 +360,7 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { // Expect to call LogManager.truncateTo exactly once val topicPartitionObj = new TopicPartition(topic, topicPartition) - val mockLogMgr: LogManager = mock(classOf[LogManager]) + val mockLogMgr = mock(classOf[LogManager]) when(mockLogMgr.liveLogDirs).thenReturn(config.logDirs.map(new File(_).getAbsoluteFile)) when(mockLogMgr.getOrCreateLog(ArgumentMatchers.eq(topicPartitionObj), ArgumentMatchers.eq(false), ArgumentMatchers.eq(false), any(), any(), any())).thenReturn(mockLog) when(mockLogMgr.getLog(topicPartitionObj, isFuture = false)).thenReturn(Some(mockLog)) @@ -388,7 +373,7 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { val aliveBrokerIds = Seq[Integer](followerBrokerId, leaderBrokerId) val aliveBrokers = aliveBrokerIds.map(brokerId => new Node(brokerId, s"host$brokerId", brokerId)) - val metadataCache: MetadataCache = mock(classOf[MetadataCache]) + val metadataCache = mock(classOf[MetadataCache]) mockGetAliveBrokerFunctions(metadataCache, aliveBrokers) when(metadataCache.getPartitionReplicaEndpoints( any[TopicPartition], any[ListenerName])). @@ -476,7 +461,7 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { @Test override def testReplicaNotAvailable(): Unit = { - def createReplicaManager(): ElasticReplicaManager = { + def createReplicaManager() = { val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect) val config = KafkaConfig.fromProps(props) val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_))) @@ -492,9 +477,7 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { alterPartitionManager = alterPartitionManager, fastFetchExecutor = MoreExecutors.newDirectExecutorService(), slowFetchExecutor = MoreExecutors.newDirectExecutorService()) { - override def getPartitionOrException(topicPartition: TopicPartition): Partition = { - throw Errors.NOT_LEADER_OR_FOLLOWER.exception() - } + override def getPartitionOrException(topicPartition: TopicPartition): Partition = throw Errors.NOT_LEADER_OR_FOLLOWER.exception() } } @@ -504,9 +487,7 @@ class ElasticReplicaManagerTest extends ReplicaManagerTest { val dir = replicaManager.logManager.liveLogDirs.head.getAbsolutePath val errors = replicaManager.alterReplicaLogDirs(Map(tp -> dir)) assertEquals(Errors.REPLICA_NOT_AVAILABLE, errors(tp)) - } finally { - replicaManager.shutdown(checkpointHW = false) - } + } finally replicaManager.shutdown(checkpointHW = false) } // Disable test preferred fetching diff --git a/docker/Dockerfile b/docker/Dockerfile index 36c5fec63c..3f08cd7157 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -13,15 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -ARG jdk_version=openjdk:17-bullseye -FROM $jdk_version +FROM docker.io/debian:bookworm + MAINTAINER AutoMQ for Apache Kafka dev@automq.com -# Set mirrors -ARG general_mirror_url="" -RUN [ -z "${general_mirror_url}" ] || (sed -i "s/deb.debian.org/${general_mirror_url}/g" /etc/apt/sources.list && sed -i "s|security.debian.org/debian-security|${general_mirror_url}/debian-security|g" /etc/apt/sources.list) -RUN apt update && apt install -y vim +RUN apt-get update && apt-get install -y ca-certificates dnsutils iputils-ping curl wget netcat-openbsd sysstat net-tools htop procps zlib1g vim less openjdk-17-jdk libjemalloc-dev ipset iproute2 && \ + apt-get clean && rm -rf /var/lib/apt/lists /var/cache/apt/archives # Do not ask for confirmations when running apt-get, etc. ENV DEBIAN_FRONTEND=noninteractive \ diff --git a/docker/README.md b/docker/README.md index 5c46ef954e..c155b26bc5 100644 --- a/docker/README.md +++ b/docker/README.md @@ -114,6 +114,29 @@ Run `pip install -r requirements.txt` to get all the requirements for running th Make sure you have docker installed with support for buildx enabled. (For pushing multi-architecture image to docker registry) +Use local code run in docker +--------------------------------------- + +- command run in project root folders + +1. generate tgz +```shell +# For example only, can be modified based on your compilation requirements +./gradlew releaseTarGz -x test -x check +``` +2. run +```shell +docker-compose -f docker/local/docker-compose.yml up -d +``` + +- After modifying your code, simply regenerate the tgz and restart the specified service. +```shell +# For example only, can be modified based on your compilation requirements +./gradlew releaseTarGz -x test -x check +# eg: restart broker +docker-compose -f docker/local/docker-compose.yml up broker1 broker2 -d --force-recreate +``` + Building image and running tests locally --------------------------------------- @@ -213,3 +236,51 @@ python generate_kafka_pr_template.py --image-type=jvm - image-type - This is the type of image that we intend to build. This will be dropdown menu type selection in the workflow. `jvm` image type is for official docker image (to be hosted on apache/kafka) as described in [KIP-975](https://cwiki.apache.org/confluence/display/KAFKA/KIP-975%3A+Docker+Image+for+Apache+Kafka). - **NOTE:** As of now [KIP-1028](https://cwiki.apache.org/confluence/display/KAFKA/KIP-1028%3A+Docker+Official+Image+for+Apache+Kafka) only aims to release JVM based Docker Official Images and not GraalVM based native Apache Kafka docker image. +AutoMQ Docker Compose Configurations +==================================== + +This directory contains Docker Compose configurations for deploying AutoMQ in different scenarios. + +Quick Start (Single Node) +------------------------- + +The main `docker-compose.yaml` in the root directory provides a simple single-node setup for quick evaluation and development: + +```bash +# From the root directory +docker-compose up -d +``` + +This configuration: +- Deploys a single AutoMQ node that acts as both controller and broker +- Includes MinIO for S3 storage +- Uses the latest bucket URI pattern (s3.data.buckets, s3.ops.buckets, s3.wal.path) +- All services run in a single Docker network + +Production-like Cluster +----------------------- + +For a more production-like setup, use the cluster configuration: + +```bash +# From the root directory +docker-compose -f docker/docker-compose-cluster.yaml up -d +``` + +This configuration: +- Deploys a 3-server cluster +- Includes MinIO for S3 storage +- Uses the latest bucket URI pattern (s3.data.buckets, s3.ops.buckets, s3.wal.path) +- All services run in a single Docker network + +Configuration Notes +------------------- + +Both configurations use the new bucket URI pattern as recommended in the AutoMQ documentation: + +- `s3.data.buckets` for data storage +- `s3.ops.buckets` for logs and metrics storage +- `s3.wal.path` for S3 WAL + +For more details, see the [AutoMQ documentation](https://www.automq.com/docs/automq/getting-started/cluster-deployment-on-linux#step-2-edit-the-cluster-configuration-template). + diff --git a/docker/automq-feat-kafka-docker.patch b/docker/automq-feat-kafka-docker.patch new file mode 100644 index 0000000000..789b0b08f8 --- /dev/null +++ b/docker/automq-feat-kafka-docker.patch @@ -0,0 +1,144 @@ +From a46795f12e3e012e6ec1ad08783f6f0428a3c034 Mon Sep 17 00:00:00 2001 +From: 1sonofqiu +Date: Wed, 10 Sep 2025 16:30:39 +0800 +Subject: [PATCH 6/6] feat(docker): add AutoMQ Kafka Docker release workflow + and update Dockerfile + +--- + docker/jvm/Dockerfile | 58 +++++++++++-------------------------------- + docker/jvm/launch | 16 ------------ + 2 files changed, 15 insertions(+), 59 deletions(-) + +diff --git a/docker/jvm/Dockerfile b/docker/jvm/Dockerfile +index 72e35e63c0..de97104271 100644 +--- a/docker/jvm/Dockerfile ++++ b/docker/jvm/Dockerfile +@@ -16,60 +16,31 @@ + # limitations under the License. + ############################################################################### + +-FROM eclipse-temurin:21-jre-alpine AS build-jsa +- +-USER root +- +-# Get kafka from https://archive.apache.org/dist/kafka and pass the url through build arguments +-ARG kafka_url +- +-COPY jsa_launch /etc/kafka/docker/jsa_launch +- +-RUN set -eux ; \ +- apk update ; \ +- apk upgrade ; \ +- apk add --no-cache wget gcompat gpg gpg-agent procps bash; \ +- mkdir opt/kafka; \ +- wget -nv -O kafka.tgz "$kafka_url"; \ +- wget -nv -O kafka.tgz.asc "$kafka_url.asc"; \ +- tar xfz kafka.tgz -C /opt/kafka --strip-components 1; \ +- wget -nv -O KEYS https://downloads.apache.org/kafka/KEYS; \ +- gpg --import KEYS; \ +- gpg --batch --verify kafka.tgz.asc kafka.tgz +- +-# Generate jsa files using dynamic CDS for kafka server start command and kafka storage format command +-RUN /etc/kafka/docker/jsa_launch +- +- +-FROM eclipse-temurin:21-jre-alpine ++FROM amazoncorretto:17-alpine + + # exposed ports + EXPOSE 9092 + + USER root + +-# Get kafka from https://archive.apache.org/dist/kafka and pass the url through build arguments ++# Get AutoMQ URL from https://github.com/AutoMQ/automq/releases and pass the url passed through build arguments + ARG kafka_url + ARG build_date + + +-LABEL org.label-schema.name="kafka" \ +- org.label-schema.description="Apache Kafka" \ ++LABEL org.label-schema.name="automq-kafka" \ ++ org.label-schema.description="AutoMQ for Kafka" \ + org.label-schema.build-date="${build_date}" \ +- org.label-schema.vcs-url="https://github.com/apache/kafka" \ +- maintainer="Apache Kafka" ++ org.label-schema.vcs-url="https://github.com/AutoMQ/automq" \ ++ maintainer="AutoMQ" + + RUN set -eux ; \ + apk update ; \ + apk upgrade ; \ +- apk add --no-cache wget gcompat gpg gpg-agent procps bash; \ ++ apk add --no-cache wget gcompat procps bash jemalloc; \ + mkdir opt/kafka; \ +- wget -nv -O kafka.tgz "$kafka_url"; \ +- wget -nv -O kafka.tgz.asc "$kafka_url.asc"; \ +- tar xfz kafka.tgz -C /opt/kafka --strip-components 1; \ +- wget -nv -O KEYS https://downloads.apache.org/kafka/KEYS; \ +- gpg --import KEYS; \ +- gpg --batch --verify kafka.tgz.asc kafka.tgz; \ ++ wget -nv -O automq.tgz "$kafka_url"; \ ++ tar xfz automq.tgz -C /opt/kafka --strip-components 1; \ + mkdir -p /var/lib/kafka/data /etc/kafka/secrets; \ + mkdir -p /etc/kafka/docker /usr/logs /mnt/shared/config; \ + adduser -h /home/appuser -D --shell /bin/bash appuser; \ +@@ -78,16 +49,17 @@ RUN set -eux ; \ + chmod -R ug+w /etc/kafka /var/lib/kafka /etc/kafka/secrets; \ + cp /opt/kafka/config/log4j.properties /etc/kafka/docker/log4j.properties; \ + cp /opt/kafka/config/tools-log4j.properties /etc/kafka/docker/tools-log4j.properties; \ +- cp /opt/kafka/config/kraft/server.properties /etc/kafka/docker/server.properties; \ +- rm kafka.tgz kafka.tgz.asc KEYS; \ +- apk del wget gpg gpg-agent; \ ++ cp /opt/kafka/config/kraft/reconfig-server.properties /etc/kafka/docker/server.properties; \ ++ rm automq.tgz; \ ++ apk del wget; \ + apk cache clean; + +-COPY --from=build-jsa kafka.jsa /opt/kafka/kafka.jsa +-COPY --from=build-jsa storage.jsa /opt/kafka/storage.jsa + COPY --chown=appuser:appuser resources/common-scripts /etc/kafka/docker + COPY --chown=appuser:appuser launch /etc/kafka/docker/launch + ++# Configure jemalloc as the memory allocator ++ENV LD_PRELOAD="/usr/lib/libjemalloc.so.2" ++ + USER appuser + + VOLUME ["/etc/kafka/secrets", "/var/lib/kafka/data", "/mnt/shared/config"] +diff --git a/docker/jvm/launch b/docker/jvm/launch +index 6c4ca1d2e0..fa0010ff03 100755 +--- a/docker/jvm/launch ++++ b/docker/jvm/launch +@@ -38,16 +38,6 @@ if [ "${KAFKA_JMX_PORT-}" ]; then + -Dcom.sun.management.jmxremote.port=$JMX_PORT" + fi + +-# Make a temp env variable to store user provided performance otps +-if [ -z "${KAFKA_JVM_PERFORMANCE_OPTS-}" ]; then +- export TEMP_KAFKA_JVM_PERFORMANCE_OPTS="" +-else +- export TEMP_KAFKA_JVM_PERFORMANCE_OPTS="$KAFKA_JVM_PERFORMANCE_OPTS" +-fi +- +-# We will first use CDS for storage to format storage +-export KAFKA_JVM_PERFORMANCE_OPTS="${KAFKA_JVM_PERFORMANCE_OPTS-} -XX:SharedArchiveFile=/opt/kafka/storage.jsa" +- + echo "===> Using provided cluster id $CLUSTER_ID ..." + + # Invoke the docker wrapper to setup property files and format storage +@@ -58,11 +48,5 @@ result=$(/opt/kafka/bin/kafka-run-class.sh kafka.docker.KafkaDockerWrapper setup + echo $result | grep -i "already formatted" || \ + { echo $result && (exit 1) } + +-# Using temp env variable to get rid of storage CDS command +-export KAFKA_JVM_PERFORMANCE_OPTS="$TEMP_KAFKA_JVM_PERFORMANCE_OPTS" +- +-# Now we will use CDS for kafka to start kafka server +-export KAFKA_JVM_PERFORMANCE_OPTS="$KAFKA_JVM_PERFORMANCE_OPTS -XX:SharedArchiveFile=/opt/kafka/kafka.jsa" +- + # Start kafka broker + exec /opt/kafka/bin/kafka-server-start.sh /opt/kafka/config/server.properties +-- +2.39.5 (Apple Git-154) + diff --git a/docker/docker-compose-cluster.yaml b/docker/docker-compose-cluster.yaml new file mode 100644 index 0000000000..23d50f6f42 --- /dev/null +++ b/docker/docker-compose-cluster.yaml @@ -0,0 +1,155 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# AutoMQ Cluster setup with MinIO for quick starts +version: "3.8" + +x-common-variables: &common-env + KAFKA_S3_ACCESS_KEY: minioadmin + KAFKA_S3_SECRET_KEY: minioadmin + KAFKA_HEAP_OPTS: -Xms1g -Xmx4g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G + # Replace CLUSTER_ID with a unique base64 UUID using "bin/kafka-storage.sh random-uuid" + CLUSTER_ID: 5XF4fHIOTfSIqkmje2KFlg + +services: + # MinIO service for S3 storage + minio: + container_name: "minio" + image: minio/minio:RELEASE.2025-05-24T17-08-30Z + environment: + MINIO_ROOT_USER: minioadmin + MINIO_ROOT_PASSWORD: minioadmin + MINIO_DOMAIN: minio + ports: + - "9000:9000" # MinIO API + - "9001:9001" # MinIO Console + command: [ "server", "/data", "--console-address", ":9001" ] + networks: + automq_net: + healthcheck: + test: [ "CMD", "curl", "-f", "http://minio:9000/minio/health/live" ] + interval: 5s + timeout: 5s + retries: 3 + + # Create needed buckets + mc: + container_name: "mc" + image: minio/mc:RELEASE.2025-05-21T01-59-54Z + depends_on: + minio: + condition: service_healthy + entrypoint: > + /bin/sh -c " + until (/usr/bin/mc alias set minio http://minio:9000 minioadmin minioadmin) do echo '...waiting...' && sleep 1; done; + /usr/bin/mc rm -r --force minio/automq-data; + /usr/bin/mc rm -r --force minio/automq-ops; + /usr/bin/mc mb minio/automq-data; + /usr/bin/mc mb minio/automq-ops; + /usr/bin/mc anonymous set public minio/automq-data; + /usr/bin/mc anonymous set public minio/automq-ops; + tail -f /dev/null + " + networks: + - automq_net + + # Three nodes for AutoMQ cluster + server1: + container_name: "automq-server1" + image: automqinc/automq:1.6.0 + stop_grace_period: 1m + environment: + <<: *common-env + command: + - bash + - -c + - | + /opt/automq/kafka/bin/kafka-server-start.sh \ + /opt/automq/kafka/config/kraft/server.properties \ + --override cluster.id=$$CLUSTER_ID \ + --override node.id=0 \ + --override controller.quorum.voters=0@server1:9093,1@server2:9093,2@server3:9093 \ + --override controller.quorum.bootstrap.servers=server1:9093,server2:9093,server3:9093 \ + --override advertised.listeners=PLAINTEXT://server1:9092 \ + --override s3.data.buckets='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.ops.buckets='1@s3://automq-ops?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.wal.path='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' + networks: + automq_net: + depends_on: + - minio + - mc + + server2: + container_name: "automq-server2" + image: automqinc/automq:1.6.0 + stop_grace_period: 1m + environment: + <<: *common-env + command: + - bash + - -c + - | + /opt/automq/kafka/bin/kafka-server-start.sh \ + /opt/automq/kafka/config/kraft/server.properties \ + --override cluster.id=$$CLUSTER_ID \ + --override node.id=1 \ + --override controller.quorum.voters=0@server1:9093,1@server2:9093,2@server3:9093 \ + --override controller.quorum.bootstrap.servers=server1:9093,server2:9093,server3:9093 \ + --override advertised.listeners=PLAINTEXT://server2:9092 \ + --override s3.data.buckets='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.ops.buckets='1@s3://automq-ops?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.wal.path='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' + networks: + automq_net: + depends_on: + - minio + - mc + + server3: + container_name: "automq-server3" + image: automqinc/automq:1.6.0 + stop_grace_period: 1m + environment: + <<: *common-env + command: + - bash + - -c + - | + /opt/automq/kafka/bin/kafka-server-start.sh \ + /opt/automq/kafka/config/kraft/server.properties \ + --override cluster.id=$$CLUSTER_ID \ + --override node.id=2 \ + --override controller.quorum.voters=0@server1:9093,1@server2:9093,2@server3:9093 \ + --override controller.quorum.bootstrap.servers=server1:9093,server2:9093,server3:9093 \ + --override advertised.listeners=PLAINTEXT://server3:9092 \ + --override s3.data.buckets='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.ops.buckets='1@s3://automq-ops?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.wal.path='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' + networks: + automq_net: + depends_on: + - minio + - mc + +networks: + automq_net: + name: automq_net + driver: bridge + ipam: + driver: default + config: + - subnet: "10.6.0.0/16" + gateway: "10.6.0.1" diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 01be4af098..10539fef14 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -13,130 +13,85 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Single-node AutoMQ setup with MinIO for quick starts version: "3.8" services: - localstack: - container_name: "${LOCALSTACK_DOCKER_NAME-localstack}" - hostname: "${LOCALSTACK_DOCKER_NAME-localstack}" - image: localstack/localstack:3.3.0 - ports: - - "4566:4566" # LocalStack Gateway - - "4510-4559:4510-4559" # external services port range + # MinIO service for S3 storage + minio: + container_name: "minio" + image: minio/minio:RELEASE.2025-05-24T17-08-30Z environment: - - DOCKER_HOST=unix:///var/run/docker.sock - volumes: - - s3_data:/var/lib/localstack - - /var/run/docker.sock:/var/run/docker.sock - # use a static ip + - MINIO_ROOT_USER=minioadmin + - MINIO_ROOT_PASSWORD=minioadmin + - MINIO_DOMAIN=minio + ports: + - "9000:9000" # MinIO API + - "9001:9001" # MinIO Console + command: [ "server", "/data", "--console-address", ":9001" ] networks: - afk_net: - ipv4_address: 10.6.0.2 + automq_net: + healthcheck: + test: [ "CMD", "curl", "-f", "http://minio:9000/minio/health/live" ] + interval: 5s + timeout: 5s + retries: 3 - # create needed buckets - aws-cli: - container_name: "${AWS_CLI_DOCKER_NAME-aws-cli}" - hostname: "${AWS_CLI_DOCKER_NAME-aws-cli}" - image: amazon/aws-cli:2.15.37 - environment: - - AWS_ACCESS_KEY_ID=test - - AWS_SECRET_ACCESS_KEY=test - - AWS_DEFAULT_REGION=us-east-1 - command: s3api create-bucket --bucket ko3 --endpoint=http://10.6.0.2:4566 + # Create needed buckets + mc: + container_name: "mc" + image: minio/mc:RELEASE.2025-05-21T01-59-54Z depends_on: - localstack: + minio: condition: service_healthy + entrypoint: > + /bin/sh -c " + until (/usr/bin/mc alias set minio http://minio:9000 minioadmin minioadmin) do echo '...waiting...' && sleep 1; done; + /usr/bin/mc rm -r --force minio/automq-data; + /usr/bin/mc rm -r --force minio/automq-ops; + /usr/bin/mc mb minio/automq-data; + /usr/bin/mc mb minio/automq-ops; + /usr/bin/mc anonymous set public minio/automq-data; + /usr/bin/mc anonymous set public minio/automq-ops; + tail -f /dev/null + " networks: - - afk_net - - controller: - container_name: "${CONTROLLER_DOCKER_NAME-controller}" - hostname: "${CONTROLLER_DOCKER_NAME-controller}" - stop_grace_period: 2m - image: automqinc/automq:1.1.0-rc0 - environment: - - KAFKA_S3_ACCESS_KEY=test - - KAFKA_S3_SECRET_KEY=test - - KAFKA_HEAP_OPTS=-Xms1g -Xmx1g -XX:MetaspaceSize=96m - command: - - bash - - -c - - | - /opt/automq/scripts/start.sh up --process.roles controller --node.id 0 --controller.quorum.voters 0@controller:9093 --s3.bucket ko3 --s3.endpoint http://10.6.0.2:4566 --s3.region us-east-1 - networks: - - afk_net - depends_on: - - localstack - - aws-cli + - automq_net - broker1: - container_name: "${BROKER1_DOCKER_NAME-broker1}" - hostname: "${BROKER1_DOCKER_NAME-broker1}" - stop_grace_period: 2m - image: automqinc/automq:1.1.0-rc0 - ports: - - "9094:9094" + # Single node with combined controller and broker roles + server1: + container_name: "automq-single-server" + image: automqinc/automq:1.6.0 + stop_grace_period: 1m environment: - - KAFKA_S3_ACCESS_KEY=test - - KAFKA_S3_SECRET_KEY=test - - KAFKA_HEAP_OPTS=-Xms1g -Xmx1g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G - - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_IN_CAPACITY=5120 - - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_OUT_CAPACITY=5120 - - KAFKA_CFG_AUTOBALANCER_REPORTER_METRICS_REPORTING_INTERVAL_MS=5000 - # override listener settings - - KAFKA_CFG_LISTENERS=PLAINTEXT://:9092,EXTERNAL://:9094 - - KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://broker1:9092,EXTERNAL://localhost:9094 - - KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT + - KAFKA_S3_ACCESS_KEY=minioadmin + - KAFKA_S3_SECRET_KEY=minioadmin + - KAFKA_HEAP_OPTS=-Xms1g -Xmx4g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G + # Replace CLUSTER_ID with a unique base64 UUID using "bin/kafka-storage.sh random-uuid" + - CLUSTER_ID=3D4fXN-yS1-vsQ8aJ_q4Mg command: - bash - -c - | - /opt/automq/scripts/start.sh up --process.roles broker --node.id 1 --controller.quorum.voters 0@controller:9093 --s3.bucket ko3 --s3.endpoint http://10.6.0.2:4566 --s3.region us-east-1 + /opt/automq/kafka/bin/kafka-server-start.sh \ + /opt/automq/kafka/config/kraft/server.properties \ + --override cluster.id=$$CLUSTER_ID \ + --override node.id=0 \ + --override controller.quorum.voters=0@server1:9093 \ + --override controller.quorum.bootstrap.servers=server1:9093 \ + --override advertised.listeners=PLAINTEXT://server1:9092 \ + --override s3.data.buckets='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.ops.buckets='1@s3://automq-ops?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.wal.path='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' networks: - - afk_net + automq_net: depends_on: - - localstack - - aws-cli - - controller - - broker2: - container_name: "${BROKER2_DOCKER_NAME-broker2}" - hostname: "${BROKER2_DOCKER_NAME-broker2}" - stop_grace_period: 2m - image: automqinc/automq:1.1.0-rc0 - ports: - - "9095:9095" - environment: - - KAFKA_S3_ACCESS_KEY=test - - KAFKA_S3_SECRET_KEY=test - - KAFKA_HEAP_OPTS=-Xms1g -Xmx1g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G - - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_IN_CAPACITY=5120 - - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_OUT_CAPACITY=5120 - - KAFKA_CFG_AUTOBALANCER_REPORTER_METRICS_REPORTING_INTERVAL_MS=5000 - # override listener settings - - KAFKA_CFG_LISTENERS=PLAINTEXT://:9092,EXTERNAL://:9095 - - KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://broker2:9092,EXTERNAL://localhost:9095 - - KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT - command: - - bash - - -c - - | - /opt/automq/scripts/start.sh up --process.roles broker --node.id 2 --controller.quorum.voters 0@controller:9093 --s3.bucket ko3 --s3.endpoint http://10.6.0.2:4566 --s3.region us-east-1 - networks: - - afk_net - depends_on: - - localstack - - aws-cli - - controller - - -volumes: - s3_data: - driver: local + - minio + - mc networks: - afk_net: - name: afk_net + automq_net: + name: automq_net driver: bridge ipam: driver: default diff --git a/docker/jvm/Dockerfile b/docker/jvm/Dockerfile index 72e35e63c0..de97104271 100644 --- a/docker/jvm/Dockerfile +++ b/docker/jvm/Dockerfile @@ -16,60 +16,31 @@ # limitations under the License. ############################################################################### -FROM eclipse-temurin:21-jre-alpine AS build-jsa - -USER root - -# Get kafka from https://archive.apache.org/dist/kafka and pass the url through build arguments -ARG kafka_url - -COPY jsa_launch /etc/kafka/docker/jsa_launch - -RUN set -eux ; \ - apk update ; \ - apk upgrade ; \ - apk add --no-cache wget gcompat gpg gpg-agent procps bash; \ - mkdir opt/kafka; \ - wget -nv -O kafka.tgz "$kafka_url"; \ - wget -nv -O kafka.tgz.asc "$kafka_url.asc"; \ - tar xfz kafka.tgz -C /opt/kafka --strip-components 1; \ - wget -nv -O KEYS https://downloads.apache.org/kafka/KEYS; \ - gpg --import KEYS; \ - gpg --batch --verify kafka.tgz.asc kafka.tgz - -# Generate jsa files using dynamic CDS for kafka server start command and kafka storage format command -RUN /etc/kafka/docker/jsa_launch - - -FROM eclipse-temurin:21-jre-alpine +FROM amazoncorretto:17-alpine # exposed ports EXPOSE 9092 USER root -# Get kafka from https://archive.apache.org/dist/kafka and pass the url through build arguments +# Get AutoMQ URL from https://github.com/AutoMQ/automq/releases and pass the url passed through build arguments ARG kafka_url ARG build_date -LABEL org.label-schema.name="kafka" \ - org.label-schema.description="Apache Kafka" \ +LABEL org.label-schema.name="automq-kafka" \ + org.label-schema.description="AutoMQ for Kafka" \ org.label-schema.build-date="${build_date}" \ - org.label-schema.vcs-url="https://github.com/apache/kafka" \ - maintainer="Apache Kafka" + org.label-schema.vcs-url="https://github.com/AutoMQ/automq" \ + maintainer="AutoMQ" RUN set -eux ; \ apk update ; \ apk upgrade ; \ - apk add --no-cache wget gcompat gpg gpg-agent procps bash; \ + apk add --no-cache wget gcompat procps bash jemalloc; \ mkdir opt/kafka; \ - wget -nv -O kafka.tgz "$kafka_url"; \ - wget -nv -O kafka.tgz.asc "$kafka_url.asc"; \ - tar xfz kafka.tgz -C /opt/kafka --strip-components 1; \ - wget -nv -O KEYS https://downloads.apache.org/kafka/KEYS; \ - gpg --import KEYS; \ - gpg --batch --verify kafka.tgz.asc kafka.tgz; \ + wget -nv -O automq.tgz "$kafka_url"; \ + tar xfz automq.tgz -C /opt/kafka --strip-components 1; \ mkdir -p /var/lib/kafka/data /etc/kafka/secrets; \ mkdir -p /etc/kafka/docker /usr/logs /mnt/shared/config; \ adduser -h /home/appuser -D --shell /bin/bash appuser; \ @@ -78,16 +49,17 @@ RUN set -eux ; \ chmod -R ug+w /etc/kafka /var/lib/kafka /etc/kafka/secrets; \ cp /opt/kafka/config/log4j.properties /etc/kafka/docker/log4j.properties; \ cp /opt/kafka/config/tools-log4j.properties /etc/kafka/docker/tools-log4j.properties; \ - cp /opt/kafka/config/kraft/server.properties /etc/kafka/docker/server.properties; \ - rm kafka.tgz kafka.tgz.asc KEYS; \ - apk del wget gpg gpg-agent; \ + cp /opt/kafka/config/kraft/reconfig-server.properties /etc/kafka/docker/server.properties; \ + rm automq.tgz; \ + apk del wget; \ apk cache clean; -COPY --from=build-jsa kafka.jsa /opt/kafka/kafka.jsa -COPY --from=build-jsa storage.jsa /opt/kafka/storage.jsa COPY --chown=appuser:appuser resources/common-scripts /etc/kafka/docker COPY --chown=appuser:appuser launch /etc/kafka/docker/launch +# Configure jemalloc as the memory allocator +ENV LD_PRELOAD="/usr/lib/libjemalloc.so.2" + USER appuser VOLUME ["/etc/kafka/secrets", "/var/lib/kafka/data", "/mnt/shared/config"] diff --git a/docker/jvm/launch b/docker/jvm/launch index 6c4ca1d2e0..fa0010ff03 100755 --- a/docker/jvm/launch +++ b/docker/jvm/launch @@ -38,16 +38,6 @@ if [ "${KAFKA_JMX_PORT-}" ]; then -Dcom.sun.management.jmxremote.port=$JMX_PORT" fi -# Make a temp env variable to store user provided performance otps -if [ -z "${KAFKA_JVM_PERFORMANCE_OPTS-}" ]; then - export TEMP_KAFKA_JVM_PERFORMANCE_OPTS="" -else - export TEMP_KAFKA_JVM_PERFORMANCE_OPTS="$KAFKA_JVM_PERFORMANCE_OPTS" -fi - -# We will first use CDS for storage to format storage -export KAFKA_JVM_PERFORMANCE_OPTS="${KAFKA_JVM_PERFORMANCE_OPTS-} -XX:SharedArchiveFile=/opt/kafka/storage.jsa" - echo "===> Using provided cluster id $CLUSTER_ID ..." # Invoke the docker wrapper to setup property files and format storage @@ -58,11 +48,5 @@ result=$(/opt/kafka/bin/kafka-run-class.sh kafka.docker.KafkaDockerWrapper setup echo $result | grep -i "already formatted" || \ { echo $result && (exit 1) } -# Using temp env variable to get rid of storage CDS command -export KAFKA_JVM_PERFORMANCE_OPTS="$TEMP_KAFKA_JVM_PERFORMANCE_OPTS" - -# Now we will use CDS for kafka to start kafka server -export KAFKA_JVM_PERFORMANCE_OPTS="$KAFKA_JVM_PERFORMANCE_OPTS -XX:SharedArchiveFile=/opt/kafka/kafka.jsa" - # Start kafka broker exec /opt/kafka/bin/kafka-server-start.sh /opt/kafka/config/server.properties diff --git a/docker/local/docker-compose.yml b/docker/local/docker-compose.yml new file mode 100644 index 0000000000..f3666f232c --- /dev/null +++ b/docker/local/docker-compose.yml @@ -0,0 +1,160 @@ + +version: "3.8" + +services: + localstack: + container_name: "${LOCALSTACK_DOCKER_NAME-localstack}" + hostname: "${LOCALSTACK_DOCKER_NAME-localstack}" + image: localstack/localstack:3.3.0 + ports: + - "4566:4566" # LocalStack Gateway + - "4510-4559:4510-4559" # external services port range + environment: + - DOCKER_HOST=unix:///var/run/docker.sock + volumes: + - s3_data:/var/lib/localstack + - /var/run/docker.sock:/var/run/docker.sock + # use a static ip + networks: + automq_net: + ipv4_address: 10.6.0.2 + + # create needed buckets + aws-cli: + container_name: "${AWS_CLI_DOCKER_NAME-aws-cli}" + hostname: "${AWS_CLI_DOCKER_NAME-aws-cli}" + image: amazon/aws-cli:2.15.37 + environment: + - AWS_ACCESS_KEY_ID=test + - AWS_SECRET_ACCESS_KEY=test + - AWS_DEFAULT_REGION=us-east-1 + command: s3api create-bucket --bucket ko3 --endpoint=http://10.6.0.2:4566 + depends_on: + localstack: + condition: service_healthy + networks: + - automq_net + +# Only comment out, do not remove. For easier future testing +# test: +# container_name: test +# hostname: test +# # image: amazoncorretto:17.0.14 +# # image: gradle:jdk17 +# image: eclipse-temurin:17-jdk-noble +# working_dir: /opt/automq +# volumes: +# - ../../core/build/distributions:/opt/volume_libs:ro +# - ../scripts:/opt/volume_scripts:ro +# command: +# - bash +# - -c +# - | +# /opt/volume_scripts/pre_start.sh && \ +# tail -f /dev/null + + controller: + container_name: "${CONTROLLER_DOCKER_NAME-controller}" + hostname: "${CONTROLLER_DOCKER_NAME-controller}" + stop_grace_period: 2m + image: eclipse-temurin:17-jdk-noble + environment: + - KAFKA_S3_ACCESS_KEY=test + - KAFKA_S3_SECRET_KEY=test + - KAFKA_HEAP_OPTS=-Xms1g -Xmx1g -XX:MetaspaceSize=96m + volumes: + - ../../core/build/distributions:/opt/volume_libs:ro + - ../scripts:/opt/volume_scripts:ro + command: + - bash + - -c + - | + /opt/volume_scripts/pre_start.sh && \ + /opt/automq/scripts/start.sh up --process.roles controller --node.id 0 --controller.quorum.voters 0@controller:9093 --s3.bucket ko3 --s3.endpoint http://10.6.0.2:4566 --s3.region us-east-1 + networks: + - automq_net + depends_on: + - localstack + - aws-cli + + broker1: + container_name: "${BROKER1_DOCKER_NAME-broker1}" + hostname: "${BROKER1_DOCKER_NAME-broker1}" + stop_grace_period: 2m + image: eclipse-temurin:17-jdk-noble + ports: + - "9094:9094" + environment: + - KAFKA_S3_ACCESS_KEY=test + - KAFKA_S3_SECRET_KEY=test + - KAFKA_HEAP_OPTS=-Xms1g -Xmx1g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G + - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_IN_CAPACITY=5120 + - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_OUT_CAPACITY=5120 + - KAFKA_CFG_AUTOBALANCER_REPORTER_METRICS_REPORTING_INTERVAL_MS=5000 + # override listener settings + - KAFKA_CFG_LISTENERS=PLAINTEXT://:9092,EXTERNAL://:9094 + - KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://broker1:9092,EXTERNAL://localhost:9094 + - KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT + volumes: + - ../../core/build/distributions:/opt/volume_libs:ro + - ../scripts:/opt/volume_scripts:ro + command: + - bash + - -c + - | + /opt/volume_scripts/pre_start.sh && \ + /opt/automq/scripts/start.sh up --process.roles broker --node.id 1 --controller.quorum.voters 0@controller:9093 --s3.bucket ko3 --s3.endpoint http://10.6.0.2:4566 --s3.region us-east-1 + networks: + - automq_net + depends_on: + - localstack + - aws-cli + - controller + + broker2: + container_name: "${BROKER2_DOCKER_NAME-broker2}" + hostname: "${BROKER2_DOCKER_NAME-broker2}" + stop_grace_period: 2m + image: eclipse-temurin:17-jdk-noble + ports: + - "9095:9095" + environment: + - KAFKA_S3_ACCESS_KEY=test + - KAFKA_S3_SECRET_KEY=test + - KAFKA_HEAP_OPTS=-Xms1g -Xmx1g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G + - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_IN_CAPACITY=5120 + - KAFKA_CFG_AUTOBALANCER_REPORTER_NETWORK_OUT_CAPACITY=5120 + - KAFKA_CFG_AUTOBALANCER_REPORTER_METRICS_REPORTING_INTERVAL_MS=5000 + # override listener settings + - KAFKA_CFG_LISTENERS=PLAINTEXT://:9092,EXTERNAL://:9095 + - KAFKA_CFG_ADVERTISED_LISTENERS=PLAINTEXT://broker2:9092,EXTERNAL://localhost:9095 + - KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP=CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT + volumes: + - ../../core/build/distributions:/opt/volume_libs:ro + - ../scripts:/opt/volume_scripts:ro + command: + - bash + - -c + - | + /opt/volume_scripts/pre_start.sh && \ + /opt/automq/scripts/start.sh up --process.roles broker --node.id 2 --controller.quorum.voters 0@controller:9093 --s3.bucket ko3 --s3.endpoint http://10.6.0.2:4566 --s3.region us-east-1 + networks: + - automq_net + depends_on: + - localstack + - aws-cli + - controller + +volumes: + s3_data: + driver: local + +networks: + automq_net: + name: automq_net + driver: bridge + ipam: + driver: default + config: + - subnet: "10.6.0.0/16" + gateway: "10.6.0.1" diff --git a/docker/scripts/pre_start.sh b/docker/scripts/pre_start.sh new file mode 100644 index 0000000000..95c9a3abd3 --- /dev/null +++ b/docker/scripts/pre_start.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash + +echo "[PreStart] mkdir" +rm -rf /opt/automq +mkdir -p /opt/kafka || exit 1 +ln -s /opt/kafka /opt/automq || exit 1 +echo "[PreStart] file" +for f in /opt/volume_libs/*.tgz; do + tar -xzf "$f" -C /opt/kafka --one-top-level=kafka --strip-components=1 --overwrite +done +cp -r /opt/volume_scripts /opt/kafka/scripts || exit 1 +find /opt/kafka -type f -name "*.sh" -exec chmod a+x {} \; +echo "[PreStart] env" +echo "export DEBIAN_FRONTEND=noninteractive" >> ~/.bashrc +echo "export AWS_DEFAULT_REGION=us-east-1" >> ~/.bashrc +echo "export KAFKA_JVM_PERFORMANCE_OPTS=\"-server -XX:+UseZGC -XX:ZCollectionInterval=5\"" >> ~/.bashrc diff --git a/docker/scripts/start.sh b/docker/scripts/start.sh index 1f33da4d42..7da5630727 100644 --- a/docker/scripts/start.sh +++ b/docker/scripts/start.sh @@ -15,6 +15,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +source ~/.bashrc + script_path="${0}" # The absolute path to the directory which this script is in. @@ -200,6 +202,7 @@ kafka_monitor_ip() { setup_value "advertised.listeners" "PLAINTEXT://${advertised_ip}:9092" "${kafka_dir}/config/kraft/${process_role}.properties" elif [[ "${process_role}" == "controller" ]]; then setup_value "listeners" "CONTROLLER://${local_private_ip}:9093" "${kafka_dir}/config/kraft/${process_role}.properties" + setup_value "advertised.listeners" "CONTROLLER://${local_private_ip}:9093" "${kafka_dir}/config/kraft/${process_role}.properties" else die "kafka_monitor_ip: unknown process role ${process_role}" fi @@ -262,9 +265,12 @@ kafka_up() { [[ -n "${s3_endpoint}" ]] || die "s3_endpoint is empty" [[ -n "${cluster_id}" ]] || cluster_id="rZdE0DjZSrqy96PXrMUZVw" + quorum_bootstrap_servers=$(echo "${quorum_voters}" | sed 's/[0-9]*@//g') + for role in "broker" "controller" "server"; do setup_value "node.id" "${node_id}" "${kafka_dir}/config/kraft/${role}.properties" - setup_value "controller.quorum.voters" "${quorum_voters}" "${kafka_dir}/config/kraft/${role}.properties" + add_or_setup_value "controller.quorum.voters" "${quorum_voters}" "${kafka_dir}/config/kraft/${role}.properties" + setup_value "controller.quorum.bootstrap.servers" "${quorum_bootstrap_servers}" "${kafka_dir}/config/kraft/${role}.properties" setup_value "s3.data.buckets" "0@s3://${s3_bucket}?region=${s3_region}&endpoint=${s3_endpoint}&authType=static" "${kafka_dir}/config/kraft/${role}.properties" setup_value "s3.ops.buckets" "0@s3://${s3_bucket}?region=${s3_region}&endpoint=${s3_endpoint}&authType=static" "${kafka_dir}/config/kraft/${role}.properties" setup_value "log.dirs" "${data_path}/kraft-${role}-logs" "${kafka_dir}/config/kraft/${role}.properties" diff --git a/docker/table_topic/README.md b/docker/table_topic/README.md new file mode 100644 index 0000000000..8ed33dcfd9 --- /dev/null +++ b/docker/table_topic/README.md @@ -0,0 +1,44 @@ +# AutoMQ Table Topic Quick Start + +This document will guide you on how to quickly start and experience the AutoMQ Table Topic feature. + +## Docker Compose Components + +This Docker Compose setup integrates the following components for a Table Topic experience: + +* **AutoMQ (`automq`)**: A single-node AutoMQ instance, supporting Table Topics and using Minio for storage. +* **Minio (`minio`)**: S3-compatible object storage for AutoMQ data and Iceberg tables. Buckets (`warehouse`, `automq-data`, `automq-ops`) are auto-created by the `mc` service. +* **Spark & Jupyter (`spark-iceberg`)**: Includes a Spark engine for Iceberg table operations and a Jupyter Notebook (accessible at `http://localhost:8888`), pre-configured for AutoMQ and Iceberg interaction. +* **Iceberg REST Catalog (`rest`)**: Metadata service for Iceberg tables, utilized by AutoMQ and Spark. +* **Schema Registry (`schema-registry`)**: Confluent Schema Registry for managing schemas used with Table Topics. + +## Usage Instructions + +1. **Start the Docker Compose Environment**: + + Execute the following command in the current directory to start all services: + + ```bash + docker-compose up -d + ``` + +2. **Access Jupyter Notebook**: + + After the services have started successfully, open `http://localhost:8888` in your web browser. + +3. **Run the Table Topic Demo Notebook**: + + In the Jupyter Notebook file browser, navigate to the `notebooks` folder, and then open the notebook named `TableTopic - Getting Started.ipynb`. + + Through this notebook, you can perform the following operations: + * **Automatic Iceberg Table Creation**: Demonstrates how Table Topic automatically creates corresponding Iceberg tables based on Kafka Topic configurations. + * **Upsert Mode**: Experience how sending messages with operation flags (Insert, Update, Delete) to a Kafka Topic synchronizes data to the Iceberg table. + * **Data Partitioning**: Understand how Table Topic partitions data for storage based on configurations. + * **Query Validation**: After each operation (insert, update, delete), query the Iceberg table using Spark SQL to verify data consistency and correctness. + * **Resource Cleanup**: Demonstrates how to delete Topics and their corresponding Iceberg tables. + +## Notes + +* Please ensure that Docker and Docker Compose are installed on your machine. +* The initial startup may take some time to download images and initialize services. +* To stop and remove all containers, use the `docker-compose down` command. diff --git a/docker/table_topic/docker-compose.yml b/docker/table_topic/docker-compose.yml new file mode 100644 index 0000000000..7a31f309d2 --- /dev/null +++ b/docker/table_topic/docker-compose.yml @@ -0,0 +1,139 @@ +version: "3" + +services: + spark-iceberg: + image: automqinc/spark-iceberg:latest + container_name: spark-iceberg + build: spark/ + networks: + iceberg_net: + depends_on: + - rest + - minio + volumes: + - ./warehouse:/home/iceberg/warehouse + - ./notebooks:/home/iceberg/notebooks/notebooks + environment: + - AWS_ACCESS_KEY_ID=admin + - AWS_SECRET_ACCESS_KEY=password + - AWS_REGION=us-east-1 + ports: + - 8888:8888 + - 8080:8080 + - 10000:10000 + - 10001:10001 + rest: + image: apache/iceberg-rest-fixture + container_name: iceberg-rest + networks: + iceberg_net: + ports: + - 8181:8181 + environment: + - AWS_ACCESS_KEY_ID=admin + - AWS_SECRET_ACCESS_KEY=password + - AWS_REGION=us-east-1 + - CATALOG_WAREHOUSE=s3://warehouse/ + - CATALOG_IO__IMPL=org.apache.iceberg.aws.s3.S3FileIO + - CATALOG_S3_ENDPOINT=http://minio:9000 + minio: + image: minio/minio:RELEASE.2025-05-24T17-08-30Z + container_name: minio + environment: + - MINIO_ROOT_USER=admin + - MINIO_ROOT_PASSWORD=password + - MINIO_DOMAIN=minio + networks: + iceberg_net: + aliases: + - warehouse.minio + ports: + - 9001:9001 + - 9000:9000 + command: ["server", "/data", "--console-address", ":9001"] + mc: + depends_on: + - minio + image: minio/mc:RELEASE.2025-05-21T01-59-54Z + container_name: mc + networks: + iceberg_net: + environment: + - AWS_ACCESS_KEY_ID=admin + - AWS_SECRET_ACCESS_KEY=password + - AWS_REGION=us-east-1 + entrypoint: | + /bin/sh -c " + until (/usr/bin/mc alias set minio http://minio:9000 admin password) do echo '...waiting...' && sleep 1; done; + /usr/bin/mc rm -r --force minio/warehouse; + /usr/bin/mc mb minio/warehouse; + /usr/bin/mc anonymous set public minio/warehouse; + /usr/bin/mc rm -r --force minio/automq-data; + /usr/bin/mc mb minio/automq-data; + /usr/bin/mc anonymous set public minio/automq-data; + /usr/bin/mc rm -r --force minio/automq-ops; + /usr/bin/mc mb minio/automq-ops; + /usr/bin/mc anonymous set public minio/automq-ops; + tail -f /dev/null + " + automq: + container_name: "automq" + image: automqinc/automq:1.6.0 + stop_grace_period: 1m + networks: + iceberg_net: + ports: + - "9092:9092" + - "9093:9093" + environment: + - KAFKA_S3_ACCESS_KEY=admin + - KAFKA_S3_SECRET_KEY=password + - KAFKA_HEAP_OPTS=-Xms1g -Xmx4g -XX:MetaspaceSize=96m -XX:MaxDirectMemorySize=1G + - CLUSTER_ID=3D4fXN-yS1-vsQ8aJ_q4Mg + command: + - bash + - -c + - | + /opt/automq/kafka/bin/kafka-server-start.sh \ + /opt/automq/kafka/config/kraft/server.properties \ + --override cluster.id=$${CLUSTER_ID} \ + --override node.id=0 \ + --override controller.quorum.voters=0@automq:9093 \ + --override controller.quorum.bootstrap.servers=automq:9093 \ + --override listeners=PLAINTEXT://:9092,CONTROLLER://:9093 \ + --override advertised.listeners=PLAINTEXT://automq:9092 \ + --override s3.data.buckets='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.ops.buckets='1@s3://automq-ops?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override s3.wal.path='0@s3://automq-data?region=us-east-1&endpoint=http://minio:9000&pathStyle=true' \ + --override automq.table.topic.catalog.type=rest \ + --override automq.table.topic.catalog.uri=http://rest:8181 \ + --override automq.table.topic.catalog.warehouse=s3://warehouse/wh/ \ + --override automq.table.topic.namespace=default \ + --override automq.table.topic.schema.registry.url=http://schema-registry:8081 + healthcheck: + test: ["CMD-SHELL", "/opt/automq/kafka/bin/kafka-broker-api-versions.sh --bootstrap-server localhost:9092 | grep -q 'automq'"] + interval: 20s + timeout: 20s + retries: 5 + start_period: 20s + depends_on: + - minio + - mc + - rest + schema-registry: + image: confluentinc/cp-schema-registry:latest + container_name: schema-registry + depends_on: + automq: + condition: service_healthy + networks: + iceberg_net: + ports: + - "8081:8081" + restart: on-failure + environment: + SCHEMA_REGISTRY_HOST_NAME: schema-registry + SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS: 'automq:9092' + SCHEMA_REGISTRY_LISTENERS: http://0.0.0.0:8081 +networks: + iceberg_net: diff --git a/docker/table_topic/spark_iceberg/Dockerfile b/docker/table_topic/spark_iceberg/Dockerfile new file mode 100644 index 0000000000..72f420bb96 --- /dev/null +++ b/docker/table_topic/spark_iceberg/Dockerfile @@ -0,0 +1,6 @@ +FROM tabulario/spark-iceberg:latest + +RUN pip install --upgrade confluent-kafka[avro] Faker pyspark + +RUN rm -rf /home/iceberg/notebooks +COPY notebooks /home/iceberg/notebooks diff --git a/docker/table_topic/spark_iceberg/notebooks/TableTopic - Getting Started.ipynb b/docker/table_topic/spark_iceberg/notebooks/TableTopic - Getting Started.ipynb new file mode 100644 index 0000000000..a768102974 --- /dev/null +++ b/docker/table_topic/spark_iceberg/notebooks/TableTopic - Getting Started.ipynb @@ -0,0 +1,357 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bff4d0e6", + "metadata": {}, + "source": [ + "# AutoMQ Table Topic Demonstration\n", + "\n", + "This notebook demonstrates the core capabilities of AutoMQ Table Topic, including automatic table creation, Upsert mode for data synchronization, and data partitioning. The workflow creates a topic with Upsert and partitioning enabled, sends one Insert (I), Update (U), and Delete (D) message, and queries the Iceberg table after each operation." + ] + }, + { + "cell_type": "markdown", + "id": "9aa48324", + "metadata": {}, + "source": [ + "## 1. Import Libraries and Define Helper Functions\n", + "\n", + "Import necessary libraries and define helper functions for key operations such as creating topics, producing messages, and querying Iceberg tables." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "345b7003", + "metadata": {}, + "outputs": [], + "source": [ + "import uuid\n", + "from confluent_kafka import Producer\n", + "from confluent_kafka.serialization import StringSerializer, SerializationContext, MessageField\n", + "from confluent_kafka.schema_registry import SchemaRegistryClient\n", + "from confluent_kafka.schema_registry.avro import AvroSerializer\n", + "from confluent_kafka.admin import AdminClient, NewTopic\n", + "from confluent_kafka.cimpl import KafkaException, KafkaError\n", + "from datetime import datetime, timezone\n", + "from faker import Faker\n", + "from pyspark.sql import SparkSession\n", + "from pyspark.sql.utils import AnalysisException\n", + "\n", + "# Configuration constants\n", + "KAFKA_BOOTSTRAP_SERVERS = 'automq:9092'\n", + "SCHEMA_REGISTRY_URL = 'http://schema-registry:8081'\n", + "TOPIC_NAME = 'web_page_view_events'\n", + "\n", + "# Initialize AdminClient and SchemaRegistryClient\n", + "admin_client_conf = {'bootstrap.servers': KAFKA_BOOTSTRAP_SERVERS}\n", + "admin_client = AdminClient(admin_client_conf)\n", + "schema_registry_conf = {'url': SCHEMA_REGISTRY_URL}\n", + "schema_registry_client = SchemaRegistryClient(schema_registry_conf)\n", + "\n", + "# Initialize SparkSession\n", + "spark = SparkSession.builder.appName(\"AutoMQ Table Topic Demo\").getOrCreate()\n", + "fake = Faker()\n", + "\n", + "# Helper function: Create a Kafka Topic\n", + "def create_topic(topic_name, num_partitions=1, replication_factor=1, config=None):\n", + " if config is None:\n", + " config = {}\n", + " topics = [NewTopic(topic_name, num_partitions=num_partitions, replication_factor=replication_factor, config=config)]\n", + " futures = admin_client.create_topics(topics, operation_timeout=30)\n", + " for topic, future in futures.items():\n", + " try:\n", + " future.result()\n", + " print(f\"Topic '{topic}' created successfully.\")\n", + " except KafkaException as e:\n", + " error = e.args[0]\n", + " if error.code() == KafkaError.TOPIC_ALREADY_EXISTS:\n", + " print(f\"Topic '{topic}' already exists.\")\n", + " else:\n", + " raise Exception(f\"Failed to create topic '{topic}': {error.str()}\")\n", + "\n", + "# Helper function: Create a Producer\n", + "def create_producer():\n", + " producer_conf = {'bootstrap.servers': KAFKA_BOOTSTRAP_SERVERS}\n", + " return Producer(producer_conf)\n", + "\n", + "# Helper function: Produce events to Kafka\n", + "def produce_events(producer, topic_name, events_data, avro_serializer, string_serializer):\n", + " for event in events_data:\n", + " try:\n", + " producer.produce(\n", + " topic=topic_name,\n", + " key=string_serializer(event.event_id),\n", + " value=avro_serializer(event, SerializationContext(topic_name, MessageField.VALUE)),\n", + " on_delivery=delivery_report\n", + " )\n", + " except Exception as e:\n", + " print(f\"Failed to produce event {event.event_id}: {e}\")\n", + " producer.poll(0)\n", + " producer.flush()\n", + " print(f\"Successfully produced {len(events_data)} event(s) to {topic_name}.\")\n", + "\n", + "# Delivery report callback for produced messages\n", + "def delivery_report(err, msg):\n", + " if err is not None:\n", + " print(f\"Message delivery failed: {err}\")\n", + " return\n", + " print(f\"Message delivered to {msg.topic()} [partition {msg.partition()}] at offset {msg.offset()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "6dca093b", + "metadata": {}, + "source": [ + "## 2. Create Topic with Upsert and Partitioning\n", + "\n", + "Create a Kafka topic with Table Topic enabled, configured for Upsert mode and partitioning. The topic uses `event_id` as the primary key, `ops` as the operation field, and partitions data by `bucket(page_url, 5)` and `hour(timestamp)`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0aa7bb63", + "metadata": {}, + "outputs": [], + "source": [ + "# Define Avro Schema with operation support\n", + "schema_str = \"\"\"\n", + "{\n", + " \"type\": \"record\",\n", + " \"name\": \"PageViewEvent\",\n", + " \"namespace\": \"com.example.events\",\n", + " \"fields\": [\n", + " {\"name\": \"event_id\", \"type\": \"string\"},\n", + " {\"name\": \"user_id\", \"type\": \"string\"},\n", + " {\"name\": \"timestamp\", \"type\": { \"type\": \"long\", \"logicalType\": \"timestamp-millis\" }},\n", + " {\"name\": \"page_url\", \"type\": \"string\"},\n", + " {\"name\": \"ip_address\", \"type\": \"string\"},\n", + " {\"name\": \"user_agent\", \"type\": \"string\"},\n", + " {\"name\": \"ops\", \"type\": \"string\"}\n", + " ]\n", + "}\n", + "\"\"\"\n", + "\n", + "# Define PageViewEvent class\n", + "class PageViewEvent:\n", + " def __init__(self, event_id, user_id, timestamp, page_url, ip_address, user_agent, ops):\n", + " self.event_id = event_id\n", + " self.user_id = user_id\n", + " self.timestamp = timestamp\n", + " self.page_url = page_url\n", + " self.ip_address = ip_address\n", + " self.user_agent = user_agent\n", + " self.ops = ops\n", + "\n", + "# Serialization function for events\n", + "def event_to_dict(event, ctx):\n", + " return {\n", + " \"event_id\": event.event_id,\n", + " \"user_id\": event.user_id,\n", + " \"timestamp\": event.timestamp,\n", + " \"page_url\": event.page_url,\n", + " \"ip_address\": event.ip_address,\n", + " \"user_agent\": event.user_agent,\n", + " \"ops\": event.ops\n", + " }\n", + "\n", + "# Create topic with Upsert and partitioning configurations\n", + "topic_config = {\n", + " 'automq.table.topic.enable': 'true',\n", + " 'automq.table.topic.commit.interval.ms': '2000',\n", + " 'automq.table.topic.schema.type': 'schema',\n", + " 'automq.table.topic.upsert.enable': 'true',\n", + " 'automq.table.topic.id.columns': '[event_id]',\n", + " 'automq.table.topic.cdc.field': 'ops',\n", + " 'automq.table.topic.partition.by': '[bucket(page_url, 5), hour(timestamp)]'\n", + "}\n", + "create_topic(TOPIC_NAME, config=topic_config)\n", + "\n", + "# Initialize serializers and producer\n", + "avro_serializer = AvroSerializer(schema_registry_client, schema_str, event_to_dict)\n", + "string_serializer = StringSerializer('utf_8')\n", + "producer = create_producer()" + ] + }, + { + "cell_type": "markdown", + "id": "11c949ad", + "metadata": {}, + "source": [ + "## 3. Insert Operation\n", + "\n", + "Produce an Insert (I) event to the topic." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2abe64e0", + "metadata": {}, + "outputs": [], + "source": [ + "event_id = str(uuid.uuid4())\n", + "current_timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)\n", + "insert_event = [PageViewEvent(event_id, fake.user_name(), current_timestamp, fake.uri_path(), fake.ipv4(), fake.user_agent(), \"I\")]\n", + "produce_events(producer, TOPIC_NAME, insert_event, avro_serializer, string_serializer)" + ] + }, + { + "cell_type": "markdown", + "id": "70a9f918", + "metadata": {}, + "source": [ + "## 4. Query After Insert\n", + "\n", + "Query the Iceberg table to verify the inserted record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ebd1757", + "metadata": {}, + "outputs": [], + "source": [ + "df = spark.read.format(\"iceberg\").load(f\"default.{TOPIC_NAME}\")\n", + "df.show()\n", + "\n", + "spark.sql(f\"SELECT file_path FROM default.{TOPIC_NAME}.files\").show(vertical=True, truncate=False)" + ] + }, + { + "cell_type": "markdown", + "id": "98922933", + "metadata": {}, + "source": [ + "## 5. Update Operation\n", + "\n", + "Produce an Update (U) event for the same `event_id` to update the record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c22f0df8", + "metadata": {}, + "outputs": [], + "source": [ + "update_event = [PageViewEvent(event_id, fake.user_name(), current_timestamp + 1000, fake.uri_path(), fake.ipv4(), fake.user_agent(), \"U\")]\n", + "produce_events(producer, TOPIC_NAME, update_event, avro_serializer, string_serializer)" + ] + }, + { + "cell_type": "markdown", + "id": "518d50a9", + "metadata": {}, + "source": [ + "## 6. Query After Update\n", + "\n", + "Query the Iceberg table to verify the updated record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3f16244a", + "metadata": {}, + "outputs": [], + "source": [ + "df = spark.read.format(\"iceberg\").load(f\"default.{TOPIC_NAME}\")\n", + "df.show()\n", + "\n", + "spark.sql(f\"SELECT file_path FROM default.{TOPIC_NAME}.files\").show(vertical=True, truncate=False)" + ] + }, + { + "cell_type": "markdown", + "id": "0f172213", + "metadata": {}, + "source": [ + "## 7. Delete Operation\n", + "\n", + "Produce a Delete (D) event for the same `event_id` to remove the record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "03b8eb79", + "metadata": {}, + "outputs": [], + "source": [ + "delete_event = [PageViewEvent(event_id, fake.user_name(), current_timestamp + 2000, fake.uri_path(), fake.ipv4(), fake.user_agent(), \"D\")]\n", + "produce_events(producer, TOPIC_NAME, delete_event, avro_serializer, string_serializer)" + ] + }, + { + "cell_type": "markdown", + "id": "89b62054", + "metadata": {}, + "source": [ + "## 8. Query After Delete\n", + "\n", + "Query the Iceberg table to verify that the record has been removed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b221c2c2", + "metadata": {}, + "outputs": [], + "source": [ + "df = spark.read.format(\"iceberg\").load(f\"default.{TOPIC_NAME}\")\n", + "df.show()\n", + "\n", + "spark.sql(f\"SELECT file_path FROM default.{TOPIC_NAME}.files\").show(vertical=True, truncate=False)" + ] + }, + { + "cell_type": "markdown", + "id": "ff13c16e", + "metadata": {}, + "source": [ + "## 9. Cleanup\n", + "\n", + "Delete the topic and drop the Iceberg table after the demonstration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15679a45", + "metadata": {}, + "outputs": [], + "source": [ + "admin_client.delete_topics([TOPIC_NAME])\n", + "spark.sql(f\"DROP TABLE default.{TOPIC_NAME}\")\n", + "print(f\"Topic '{TOPIC_NAME}' and Iceberg table deleted.\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.16" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/images/automq-architecture.png b/docs/images/automq-architecture.png deleted file mode 100644 index 38dd5059d7..0000000000 Binary files a/docs/images/automq-architecture.png and /dev/null differ diff --git a/docs/images/automq-kafka-compare.png b/docs/images/automq-kafka-compare.png deleted file mode 100644 index 24b6884e73..0000000000 Binary files a/docs/images/automq-kafka-compare.png and /dev/null differ diff --git a/docs/images/automq_dashboard.gif b/docs/images/automq_dashboard.gif deleted file mode 100644 index a8c9402448..0000000000 Binary files a/docs/images/automq_dashboard.gif and /dev/null differ diff --git a/docs/images/automq_dashboard.jpeg b/docs/images/automq_dashboard.jpeg deleted file mode 100644 index ffcf6e8b04..0000000000 Binary files a/docs/images/automq_dashboard.jpeg and /dev/null differ diff --git a/docs/images/automq_s3stream_architecture.gif b/docs/images/automq_s3stream_architecture.gif deleted file mode 100644 index 26c9299c60..0000000000 Binary files a/docs/images/automq_s3stream_architecture.gif and /dev/null differ diff --git a/docs/images/automq_simple_arch.png b/docs/images/automq_simple_arch.png new file mode 100644 index 0000000000..5eaa445ecd Binary files /dev/null and b/docs/images/automq_simple_arch.png differ diff --git a/docs/images/automq_vs_kafka.gif b/docs/images/automq_vs_kafka.gif deleted file mode 100644 index b6164aa74b..0000000000 Binary files a/docs/images/automq_vs_kafka.gif and /dev/null differ diff --git a/docs/images/automq_wal_architecture.gif b/docs/images/automq_wal_architecture.gif deleted file mode 100644 index f537a8d58b..0000000000 Binary files a/docs/images/automq_wal_architecture.gif and /dev/null differ diff --git a/docs/images/banner-readme.jpeg b/docs/images/banner-readme.jpeg deleted file mode 100644 index ab0fee0cc0..0000000000 Binary files a/docs/images/banner-readme.jpeg and /dev/null differ diff --git a/examples/src/main/java/kafka/examples/TransactionProducer.java b/examples/src/main/java/kafka/examples/TransactionProducer.java index 71fb318c92..16a9bde173 100644 --- a/examples/src/main/java/kafka/examples/TransactionProducer.java +++ b/examples/src/main/java/kafka/examples/TransactionProducer.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package kafka.examples; diff --git a/gradle/dependencies.gradle b/gradle/dependencies.gradle index 6f24e5ecb3..398bfbefaf 100644 --- a/gradle/dependencies.gradle +++ b/gradle/dependencies.gradle @@ -103,7 +103,6 @@ versions += [ jetty: "9.4.54.v20240208", jersey: "2.39.1", jline: "3.25.1", - jline: "3.25.1", jmh: "1.37", hamcrest: "2.2", scalaLogging: "3.9.5", @@ -155,6 +154,7 @@ versions += [ scalaJava8Compat : "1.0.2", scoverage: "2.0.11", slf4j: "1.7.36", + jclOverSlf4j: "1.7.36", snappy: "1.1.10.5", spotbugs: "4.8.0", zinc: "1.9.2", @@ -169,12 +169,17 @@ versions += [ opentelemetrySDKAlpha: "1.40.0-alpha", opentelemetryInstrument: "2.6.0-alpha", oshi: "6.4.7", - awsSdk:"2.26.10", + awsSdk:"2.29.26", bucket4j:"8.5.0", jna:"5.2.0", guava:"32.0.1-jre", hdrHistogram:"2.1.12", - nettyTcnativeBoringSsl: "2.0.65.Final", + nettyTcnativeBoringSsl: "2.0.69.Final", + avro: "1.11.4", + confluentSchema: "7.8.0", + iceberg: "1.6.1", + wire: "4.9.1", + oshi: "6.8.1", // AutoMQ inject end junitPlatform: "1.10.2" @@ -272,6 +277,7 @@ libs += [ scalaReflect: "org.scala-lang:scala-reflect:$versions.scala", slf4jApi: "org.slf4j:slf4j-api:$versions.slf4j", slf4jReload4j: "org.slf4j:slf4j-reload4j:$versions.slf4j", + jclOverSlf4j: "org.slf4j:jcl-over-slf4j:$versions.jclOverSlf4j", slf4jBridge: "org.slf4j:jul-to-slf4j:$versions.slf4j", snappy: "org.xerial.snappy:snappy-java:$versions.snappy", swaggerAnnotations: "io.swagger.core.v3:swagger-annotations:$swaggerVersion", @@ -300,5 +306,7 @@ libs += [ jna: "net.java.dev.jna:jna:$versions.jna", guava: "com.google.guava:guava:$versions.guava", hdrHistogram: "org.hdrhistogram:HdrHistogram:$versions.hdrHistogram", + kafkaAvroSerializer: "io.confluent:kafka-avro-serializer:$versions.confluentSchema", spotbugsAnnotations: "com.github.spotbugs:spotbugs-annotations:$versions.spotbugs", + oshi: "com.github.oshi:oshi-core:$versions.oshi", ] diff --git a/gradle/spotbugs-exclude.xml b/gradle/spotbugs-exclude.xml index 310d9902d9..d4062a88b5 100644 --- a/gradle/spotbugs-exclude.xml +++ b/gradle/spotbugs-exclude.xml @@ -601,7 +601,8 @@ For a detailed description of spotbugs bug categories, see https://spotbugs.read - + + + + + + + + + + + diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinator.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinator.java index 2df6f1136b..0033100b70 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinator.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinator.java @@ -17,6 +17,8 @@ package org.apache.kafka.coordinator.group; import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.message.AutomqUpdateGroupRequestData; +import org.apache.kafka.common.message.AutomqUpdateGroupResponseData; import org.apache.kafka.common.message.ConsumerGroupDescribeResponseData; import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData; import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData; @@ -413,4 +415,22 @@ void onNewMetadataImage( * Shutdown the group coordinator. */ void shutdown(); + + // AutoMQ injection start + /** + * Update consumer groups + * + * @param context The coordinator request context. + * @param request The AutomqUpdateGroupRequestData data. + * @param bufferSupplier The buffer supplier tight to the request thread. + * + * @return A future yielding the response. + * The error code(s) of the response are set to indicate the error(s) occurred during the execution. + */ + CompletableFuture updateGroup( + RequestContext context, + AutomqUpdateGroupRequestData request, + BufferSupplier bufferSupplier + ); + // AutoMQ injection end } diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java index e132e85ec2..edc28db589 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorService.java @@ -21,6 +21,8 @@ import org.apache.kafka.common.config.TopicConfig; import org.apache.kafka.common.errors.NotCoordinatorException; import org.apache.kafka.common.internals.Topic; +import org.apache.kafka.common.message.AutomqUpdateGroupRequestData; +import org.apache.kafka.common.message.AutomqUpdateGroupResponseData; import org.apache.kafka.common.message.ConsumerGroupDescribeResponseData; import org.apache.kafka.common.message.ConsumerGroupHeartbeatRequestData; import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData; @@ -1118,6 +1120,12 @@ public void shutdown() { log.info("Shutdown complete."); } + @Override + public CompletableFuture updateGroup(RequestContext context, + AutomqUpdateGroupRequestData request, BufferSupplier bufferSupplier) { + return FutureUtils.failedFuture(new UnsupportedOperationException()); + } + private static boolean isGroupIdNotEmpty(String groupId) { return groupId != null && !groupId.isEmpty(); } diff --git a/licenses/LICENSE-binary b/licenses/LICENSE-binary index 01a2b03025..b0355b588e 100644 --- a/licenses/LICENSE-binary +++ b/licenses/LICENSE-binary @@ -1,7 +1,3 @@ -The components built from this repository are licensed under -the Business Source License 1.1 (BSL-1.1) -or the Apache License Version 2.0 if the source code is inherited from Apache projects. -------------------------------------------------------------------------------- This project bundles some components that are also licensed under the Apache License Version 2.0: diff --git a/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java b/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java index e54615cc77..8be03d2280 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java @@ -20,6 +20,7 @@ import org.apache.kafka.common.message.BrokerHeartbeatRequestData; import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.common.utils.Time; +import org.apache.kafka.controller.stream.NodeState; import org.apache.kafka.controller.stream.OverloadCircuitBreaker; import org.apache.kafka.metadata.placement.UsableBroker; @@ -83,6 +84,16 @@ static class BrokerHeartbeatState { */ private long controlledShutdownOffset; + // AutoMQ inject start + /** + * The last time the broker was controlled shutdown, in monotonic nanoseconds, or 0 + * if the broker has never been controlled shutdown since the most recent start. + * It will be updated on receiving a broker heartbeat with controlled shutdown request. + * It will be reset to 0 when the broker is active again. + */ + private long lastControlledShutdownNs; + // AutoMQ inject end + /** * The previous entry in the unfenced list, or null if the broker is not in that list. */ @@ -100,6 +111,9 @@ static class BrokerHeartbeatState { this.next = null; this.metadataOffset = -1; this.controlledShutdownOffset = -1; + // AutoMQ inject start + this.lastControlledShutdownNs = 0; + // AutoMQ inject end } /** @@ -122,6 +136,12 @@ boolean fenced() { boolean shuttingDown() { return controlledShutdownOffset >= 0; } + + // AutoMQ inject start + long lastControlledShutdownNs() { + return lastControlledShutdownNs; + } + // AutoMQ inject end } static class MetadataOffsetComparator implements Comparator { @@ -441,6 +461,9 @@ void maybeUpdateControlledShutdownOffset(int brokerId, long controlledShutDownOf throw new RuntimeException("Fenced brokers cannot enter controlled shutdown."); } active.remove(broker); + // AutoMQ inject start + broker.lastControlledShutdownNs = time.nanoseconds(); + // AutoMQ inject end if (broker.controlledShutdownOffset < 0) { broker.controlledShutdownOffset = controlledShutDownOffset; log.debug("Updated the controlled shutdown offset for broker {} to {}.", @@ -489,6 +512,24 @@ Iterator usableBrokers( } // AutoMQ inject start + public NodeState brokerState(int brokerId, long shutdownTimeoutNs) { + BrokerHeartbeatState broker = brokers.get(brokerId); + if (broker == null) { + return NodeState.UNKNOWN; + } + if (broker.shuttingDown()) { + return NodeState.CONTROLLED_SHUTDOWN; + } + if (broker.fenced()) { + if (broker.lastControlledShutdownNs() + shutdownTimeoutNs > time.nanoseconds()) { + // The broker is still in controlled shutdown. + return NodeState.CONTROLLED_SHUTDOWN; + } + return NodeState.FENCED; + } + return NodeState.ACTIVE; + } + long nextCheckTimeNs() { if (overloadCircuitBreaker.isOverload()) { return Long.MAX_VALUE; diff --git a/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java index 086f2f3ce7..dbaf5c5045 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java @@ -27,6 +27,8 @@ import org.apache.kafka.common.errors.UnsupportedVersionException; import org.apache.kafka.common.message.BrokerRegistrationRequestData; import org.apache.kafka.common.message.ControllerRegistrationRequestData; +import org.apache.kafka.common.message.GetKVsRequestData; +import org.apache.kafka.common.message.PutKVsRequestData; import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord; import org.apache.kafka.common.metadata.FenceBrokerRecord; import org.apache.kafka.common.metadata.RegisterBrokerRecord; @@ -39,6 +41,7 @@ import org.apache.kafka.common.protocol.ApiMessage; import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.common.utils.Time; +import org.apache.kafka.controller.stream.KVControlManager; import org.apache.kafka.metadata.BrokerRegistration; import org.apache.kafka.metadata.BrokerRegistrationFencingChange; import org.apache.kafka.metadata.BrokerRegistrationInControlledShutdownChange; @@ -59,8 +62,10 @@ import org.slf4j.Logger; +import java.nio.ByteBuffer; import java.util.AbstractMap; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; @@ -100,11 +105,17 @@ static class Builder { // AutoMQ for Kafka inject start private List quorumVoters; + private KVControlManager kvControlManager; Builder setQuorumVoters(List quorumVoters) { this.quorumVoters = quorumVoters; return this; } + + Builder setKVControlManager(KVControlManager kvControlManager) { + this.kvControlManager = kvControlManager; + return this; + } // AutoMQ for Kafka inject end Builder setLogContext(LogContext logContext) { @@ -180,7 +191,10 @@ ClusterControlManager build() { featureControl, zkMigrationEnabled, brokerUncleanShutdownHandler, - quorumVoters + // AutoMQ inject start + quorumVoters, + kvControlManager + // AutoMQ inject end ); } } @@ -292,6 +306,12 @@ boolean check() { * The real next available node id is generally one greater than this value. */ private AtomicInteger nextNodeId = new AtomicInteger(-1); + + /** + * A set of node IDs that have been unregistered and can be reused for new node assignments. + */ + private final KVControlManager kvControlManager; + private static final String REUSABLE_NODE_IDS_KEY = "__automq_reusable_node_ids/"; // AutoMQ for Kafka inject end private ClusterControlManager( @@ -304,7 +324,10 @@ private ClusterControlManager( FeatureControlManager featureControl, boolean zkMigrationEnabled, BrokerUncleanShutdownHandler brokerUncleanShutdownHandler, - List quorumVoters + // AutoMQ inject start + List quorumVoters, + KVControlManager kvControlManager + // AutoMQ inject end ) { this.logContext = logContext; this.clusterId = clusterId; @@ -323,6 +346,7 @@ private ClusterControlManager( this.brokerUncleanShutdownHandler = brokerUncleanShutdownHandler; // AutoMQ for Kafka inject start this.maxControllerId = QuorumConfig.parseVoterConnections(quorumVoters).keySet().stream().max(Integer::compareTo).orElse(0); + this.kvControlManager = kvControlManager; // AutoMQ for Kafka inject end } @@ -369,16 +393,73 @@ boolean zkRegistrationAllowed() { // AutoMQ for Kafka inject start public ControllerResult getNextNodeId() { - int maxBrokerId = brokerRegistrations.keySet().stream().max(Integer::compareTo).orElse(0); - int maxNodeId = Math.max(maxBrokerId, maxControllerId); - int nextId = this.nextNodeId.accumulateAndGet(maxNodeId, (x, y) -> Math.max(x, y) + 1); - // Let the broker's nodeId start from 1000 to easily distinguish broker and controller. - nextId = Math.max(nextId, 1000); - UpdateNextNodeIdRecord record = new UpdateNextNodeIdRecord().setNodeId(nextId); + int nextId; + Set reusableNodeIds = getReusableNodeIds(); + if (!reusableNodeIds.isEmpty()) { + Iterator iterator = reusableNodeIds.iterator(); + nextId = iterator.next(); + // we simply remove the id from reusable id set because we're unable to determine if the id + // will finally be used. + iterator.remove(); + return ControllerResult.atomicOf(putReusableNodeIds(reusableNodeIds), nextId); + } else { + int maxBrokerId = brokerRegistrations.keySet().stream().max(Integer::compareTo).orElse(0); + int maxNodeId = Math.max(maxBrokerId, maxControllerId); + nextId = this.nextNodeId.accumulateAndGet(maxNodeId, (x, y) -> Math.max(x, y) + 1); + // Let the broker's nodeId start from 1000 to easily distinguish broker and controller. + nextId = Math.max(nextId, 1000); + UpdateNextNodeIdRecord record = new UpdateNextNodeIdRecord().setNodeId(nextId); - List records = new ArrayList<>(); - records.add(new ApiMessageAndVersion(record, (short) 0)); - return ControllerResult.atomicOf(records, nextId); + List records = new ArrayList<>(); + records.add(new ApiMessageAndVersion(record, (short) 0)); + return ControllerResult.atomicOf(records, nextId); + } + } + + Set getReusableNodeIds() { + return deserializeReusableNodeIds(kvControlManager.getKV( + new GetKVsRequestData.GetKVRequest().setKey(REUSABLE_NODE_IDS_KEY)).value()); + } + + List putReusableNodeIds(Set reusableNodeIds) { + return kvControlManager.putKV(new PutKVsRequestData.PutKVRequest() + .setKey(REUSABLE_NODE_IDS_KEY) + .setValue(serializeReusableNodeIds(reusableNodeIds)) + .setOverwrite(true)) + .records(); + } + + private Set deserializeReusableNodeIds(byte[] value) { + if (value == null) { + return new HashSet<>(); + } + ByteBuffer buffer = ByteBuffer.wrap(value); + Set reusableNodeIds = new HashSet<>(); + while (buffer.hasRemaining()) { + reusableNodeIds.add(buffer.getInt()); + } + return reusableNodeIds; + } + + private byte[] serializeReusableNodeIds(Set reusableNodeIds) { + ByteBuffer buffer = ByteBuffer.allocate(reusableNodeIds.size() * Integer.BYTES); + reusableNodeIds.forEach(buffer::putInt); + return buffer.array(); + } + + public List registerBrokerRecords(int brokerId) { + Set reusableNodeIds = getReusableNodeIds(); + if (reusableNodeIds.contains(brokerId)) { + reusableNodeIds.remove(brokerId); + return putReusableNodeIds(reusableNodeIds); + } + return Collections.emptyList(); + } + + public List unRegisterBrokerRecords(int brokerId) { + Set reusableNodeIds = getReusableNodeIds(); + reusableNodeIds.add(brokerId); + return putReusableNodeIds(reusableNodeIds); } // AutoMQ for Kafka inject end @@ -496,6 +577,10 @@ public ControllerResult registerBroker( } heartbeatManager.register(brokerId, record.fenced()); + // AutoMQ for Kafka inject start + records.addAll(registerBrokerRecords(brokerId)); + // AutoMQ for Kafka inject end + return ControllerResult.atomicOf(records, new BrokerRegistrationReply(record.brokerEpoch())); } @@ -583,6 +668,7 @@ public void replay(RegisterBrokerRecord record, long offset) { if (prevRegistration != null) heartbeatManager.remove(brokerId); heartbeatManager.register(brokerId, record.fenced()); } + if (prevRegistration == null) { log.info("Replayed initial RegisterBrokerRecord for broker {}: {}", record.brokerId(), record); } else if (prevRegistration.incarnationId().equals(record.incarnationId())) { @@ -608,6 +694,7 @@ public void replay(UnregisterBrokerRecord record) { if (heartbeatManager != null) heartbeatManager.remove(brokerId); updateDirectories(brokerId, registration.directories(), null); brokerRegistrations.remove(brokerId); + // AutoMQ injection end log.info("Replayed {}", record); } } @@ -880,5 +967,9 @@ public List getActiveBrokers() { .filter(b -> isActive(b.id())) .collect(Collectors.toList()); } + + public BrokerHeartbeatManager getHeartbeatManager() { + return heartbeatManager; + } // AutoMQ inject end } diff --git a/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java b/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java index 3f3ef07340..cd94e6a566 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java +++ b/metadata/src/main/java/org/apache/kafka/controller/QuorumController.java @@ -145,9 +145,10 @@ import org.apache.kafka.controller.errors.ControllerExceptions; import org.apache.kafka.controller.errors.EventHandlerExceptionInfo; import org.apache.kafka.controller.metrics.QuorumControllerMetrics; -import org.apache.kafka.controller.stream.DefaultNodeRuntimeInfoGetter; +import org.apache.kafka.controller.stream.DefaultNodeRuntimeInfoManager; import org.apache.kafka.controller.stream.KVControlManager; import org.apache.kafka.controller.stream.NodeControlManager; +import org.apache.kafka.controller.stream.RouterChannelEpochControlManager; import org.apache.kafka.controller.stream.S3ObjectControlManager; import org.apache.kafka.controller.stream.StreamClient; import org.apache.kafka.controller.stream.StreamControlManager; @@ -664,7 +665,7 @@ public void run() throws Exception { } finally { long processTime = NANOSECONDS.toMicros(time.nanoseconds() - startProcessingTimeNs.getAsLong()); if (processTime > EventQueue.Event.EVENT_PROCESS_TIME_THRESHOLD_MICROSECOND) { - log.error("Controller took {} µs to process control event: {}", processTime, name); + log.warn("Controller took {} µs to process control event: {}", processTime, name); } } handleEventEnd(this.toString(), startProcessingTimeNs.getAsLong()); @@ -722,7 +723,7 @@ public void run() throws Exception { } finally { long processTime = NANOSECONDS.toMicros(time.nanoseconds() - startProcessingTimeNs.getAsLong()); if (processTime > EventQueue.Event.EVENT_PROCESS_TIME_THRESHOLD_MICROSECOND) { - log.error("Controller took {} µs to process read event: {}", processTime, name); + log.warn("Controller took {} µs to process read event: {}", processTime, name); } } handleEventEnd(this.toString(), startProcessingTimeNs.getAsLong()); @@ -760,6 +761,10 @@ public SnapshotRegistry snapshotRegistry() { return snapshotRegistry; } + public StreamControlManager streamControlManager() { + return streamControlManager; + } + public NodeControlManager nodeControlManager() { return nodeControlManager; } @@ -899,7 +904,7 @@ public void run() throws Exception { } finally { long processTime = NANOSECONDS.toMicros(time.nanoseconds() - startProcessingTimeNs.getAsLong()); if (processTime > EventQueue.Event.EVENT_PROCESS_TIME_THRESHOLD_MICROSECOND) { - log.error("Controller took {} µs to process write event: {}", processTime, name); + log.warn("Controller took {} µs to process write event: {}", processTime, name); } } @@ -1737,6 +1742,7 @@ private void replay(ApiMessage message, Optional snapshotId, lon kvControlManager.replay(record); topicDeletionManager.replay(record); nodeControlManager.replay(record); + routerChannelEpochControlManager.replay(record); break; } case REMOVE_KVRECORD: { @@ -1744,6 +1750,7 @@ private void replay(ApiMessage message, Optional snapshotId, lon kvControlManager.replay(record); topicDeletionManager.replay(record); nodeControlManager.replay(record); + routerChannelEpochControlManager.replay(record); break; } case UPDATE_NEXT_NODE_ID_RECORD: @@ -1994,6 +2001,11 @@ private enum ImbalanceSchedule { */ private final NodeControlManager nodeControlManager; + /** + * Manage the router channel epoch; + */ + private final RouterChannelEpochControlManager routerChannelEpochControlManager; + private final QuorumControllerExtension extension; // AutoMQ for Kafka inject end @@ -2047,6 +2059,9 @@ private QuorumController( this.time = time; this.controllerMetrics = controllerMetrics; this.snapshotRegistry = new SnapshotRegistry(logContext); + // AutoMQ for Kafka inject start + this.kvControlManager = new KVControlManager(snapshotRegistry, logContext); + // AutoMQ for Kafka inject end this.deferredEventQueue = new DeferredEventQueue(logContext); this.deferredUnstableEventQueue = new DeferredEventQueue(logContext); this.offsetControl = new OffsetControlManager.Builder(). @@ -2094,6 +2109,7 @@ private QuorumController( setZkMigrationEnabled(zkMigrationEnabled). // AutoMQ for Kafka inject start setQuorumVoters(quorumVoters). + setKVControlManager(kvControlManager). // AutoMQ for Kafka inject end setBrokerUncleanShutdownHandler(this::handleUncleanBrokerShutdown). build(); @@ -2156,9 +2172,9 @@ private QuorumController( featureControl::autoMQVersion, time); this.streamControlManager = new StreamControlManager(this, snapshotRegistry, logContext, this.s3ObjectControlManager, clusterControl, featureControl, replicationControl); - this.kvControlManager = new KVControlManager(snapshotRegistry, logContext); this.topicDeletionManager = new TopicDeletionManager(snapshotRegistry, this, streamControlManager, kvControlManager); - this.nodeControlManager = new NodeControlManager(snapshotRegistry, new DefaultNodeRuntimeInfoGetter(clusterControl, streamControlManager)); + this.nodeControlManager = new NodeControlManager(snapshotRegistry, new DefaultNodeRuntimeInfoManager(clusterControl, streamControlManager)); + this.routerChannelEpochControlManager = new RouterChannelEpochControlManager(snapshotRegistry, this, nodeControlManager, time); this.extension = extension.apply(this); // set the nodeControlManager here to avoid circular dependency diff --git a/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java index 7104ea1945..51ed11b338 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java @@ -1486,6 +1486,7 @@ void handleBrokerUnregistered(int brokerId, long brokerEpoch, (short) 0)); // AutoMQ for Kafka inject start records.add(nodeControlManager.unregisterNodeRecord(brokerId)); + records.addAll(clusterControl.unRegisterBrokerRecords(brokerId)); // AutoMQ for Kafka inject end } @@ -1543,7 +1544,7 @@ void handleBrokerInControlledShutdown(int brokerId, long brokerEpoch, List records, Iterator iterator) { - generateLeaderAndIsrUpdates(context, brokerToRemove, brokerToAdd, brokerWithUncleanShutdown, records, iterator, false); + generateLeaderAndIsrUpdates0(context, brokerToRemove, brokerToAdd, brokerWithUncleanShutdown, records, iterator); } - /** - * Iterate over a sequence of partitions and generate ISR changes and/or leader - * changes if necessary. - * - * @param context A human-readable context string used in log4j logging. - * @param brokerToRemove NO_LEADER if no broker is being removed; the ID of the - * broker to remove from the ISR and leadership, otherwise. - * @param brokerToAdd NO_LEADER if no broker is being added; the ID of the - * broker which is now eligible to be a leader, otherwise. - * @param records A list of records which we will append to. - * @param iterator The iterator containing the partitions to examine. - * @param fencing Whether to fence the provided partitions. That is to say, - * set their leader to {@link org.apache.kafka.metadata.LeaderConstants#NO_LEADER} - * temporarily. It aims to ensure that the partitions should be firstly closed and - * then be re-opened. In case that the original broker is out of communication and - * then fail to touch re-elections, The partitions are scheduled to be re-elected. - */ - void generateLeaderAndIsrUpdates(String context, + void generateLeaderAndIsrUpdates0(String context, int brokerToRemove, int brokerToAdd, int brokerWithUncleanShutdown, List records, - Iterator iterator, - boolean fencing) { + Iterator iterator) { int oldSize = records.size(); // If the caller passed a valid broker ID for brokerToAdd, rather than passing @@ -2118,8 +2101,9 @@ void generateLeaderAndIsrUpdates(String context, // the ISR. // AutoMQ for Kafka inject start - IntPredicate isAcceptableLeader = fencing ? r -> false : - r -> (r != brokerToRemove) && (r == brokerToAdd || clusterControl.isActive(r)); + // We should set up set leader after the partition is opened in the new broker to avoid client fast retry. + // When the partition is opened in a new broker, the new broker will try to elect leader for the partition. + IntPredicate isAcceptableLeader = n -> false; BrokerRegistration brokerRegistrationToRemove = clusterControl.brokerRegistrations().get(brokerToRemove); PartitionLeaderSelector partitionLeaderSelector = null; @@ -2170,10 +2154,6 @@ void generateLeaderAndIsrUpdates(String context, .select(new TopicPartition(topic.name(), topicIdPart.partitionId())) .ifPresent(builder::setTargetNode); } - if (fencing) { - TopicPartition topicPartition = new TopicPartition(topic.name(), topicIdPart.partitionId()); - addPartitionToReElectTimeouts(topicPartition); - } } else { builder.setTargetIsr(Replicas.toList( Replicas.copyWithout(partition.isr, new int[] {brokerToRemove, brokerWithUncleanShutdown}))); diff --git a/metadata/src/main/java/org/apache/kafka/controller/automq/utils/AvroUtils.java b/metadata/src/main/java/org/apache/kafka/controller/automq/utils/AvroUtils.java new file mode 100644 index 0000000000..ae2e5307c1 --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/controller/automq/utils/AvroUtils.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.controller.automq.utils; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.DatumReader; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.Decoder; +import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.Encoder; +import org.apache.avro.io.EncoderFactory; +import org.apache.avro.specific.SpecificDatumReader; +import org.apache.avro.specific.SpecificDatumWriter; + +import java.io.IOException; +import java.util.function.Function; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.ByteBufInputStream; +import io.netty.buffer.ByteBufOutputStream; +import io.netty.buffer.Unpooled; + +public class AvroUtils { + + public static ByteBuf encode(GenericRecord record, short version) throws IOException { + try (ByteBufOutputStream out = new ByteBufOutputStream(Unpooled.buffer())) { + out.writeShort(version); + DatumWriter writer = new SpecificDatumWriter<>(record.getSchema()); + Encoder encoder = EncoderFactory.get().binaryEncoder(out, null); + writer.write(record, encoder); + encoder.flush(); + return out.buffer(); + } + } + + public static GenericRecord decode(ByteBuf buf, Function schemaGetter) throws IOException { + buf = buf.slice(); + short version = buf.readShort(); + Schema schema = schemaGetter.apply(version); + DatumReader reader = new SpecificDatumReader<>(schema); + Decoder decoder = DecoderFactory.get().binaryDecoder(new ByteBufInputStream(buf), null); + return reader.read(null, decoder); + } + +} diff --git a/metadata/src/main/java/org/apache/kafka/controller/es/ClusterStats.java b/metadata/src/main/java/org/apache/kafka/controller/es/ClusterStats.java index efa744885c..6181fb447e 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/es/ClusterStats.java +++ b/metadata/src/main/java/org/apache/kafka/controller/es/ClusterStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.es; diff --git a/metadata/src/main/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelector.java b/metadata/src/main/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelector.java index fdced1c9a6..a37bea69a3 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelector.java +++ b/metadata/src/main/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelector.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.es; diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoGetter.java b/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoGetter.java deleted file mode 100644 index e71d391afa..0000000000 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoGetter.java +++ /dev/null @@ -1,45 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package org.apache.kafka.controller.stream; - -import org.apache.kafka.controller.ClusterControlManager; -import org.apache.kafka.metadata.BrokerRegistration; - -public class DefaultNodeRuntimeInfoGetter implements NodeRuntimeInfoGetter { - private final ClusterControlManager clusterControlManager; - private final StreamControlManager streamControlManager; - - public DefaultNodeRuntimeInfoGetter(ClusterControlManager clusterControlManager, StreamControlManager streamControlManager) { - this.clusterControlManager = clusterControlManager; - this.streamControlManager = streamControlManager; - } - - @Override - public NodeState state(int nodeId) { - BrokerRegistration brokerRegistration = clusterControlManager.registration(nodeId); - if (brokerRegistration == null) { - return NodeState.UNKNOWN; - } - if (brokerRegistration.fenced()) { - return NodeState.FENCED; - } - if (brokerRegistration.inControlledShutdown()) { - return NodeState.CONTROLLED_SHUTDOWN; - } - return NodeState.ACTIVE; - } - - @Override - public boolean hasOpeningStreams(int nodeId) { - return streamControlManager.hasOpeningStreams(nodeId); - } -} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoManager.java new file mode 100644 index 0000000000..d50e5dcd0a --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/DefaultNodeRuntimeInfoManager.java @@ -0,0 +1,62 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.controller.stream; + +import org.apache.kafka.controller.BrokerHeartbeatManager; +import org.apache.kafka.controller.ClusterControlManager; + +import java.util.concurrent.TimeUnit; + +public class DefaultNodeRuntimeInfoManager implements NodeRuntimeInfoManager { + private static final long SHUTDOWN_TIMEOUT_NS = TimeUnit.SECONDS.toNanos(60); + + private final ClusterControlManager clusterControlManager; + private final StreamControlManager streamControlManager; + + public DefaultNodeRuntimeInfoManager(ClusterControlManager clusterControlManager, StreamControlManager streamControlManager) { + this.clusterControlManager = clusterControlManager; + this.streamControlManager = streamControlManager; + } + + @Override + public NodeState state(int nodeId) { + BrokerHeartbeatManager brokerHeartbeatManager = clusterControlManager.getHeartbeatManager(); + if (null == brokerHeartbeatManager) { + // This controller is not the active controller, so we don't have the heartbeat manager. + return NodeState.UNKNOWN; + } + return brokerHeartbeatManager.brokerState(nodeId, SHUTDOWN_TIMEOUT_NS); + } + + @Override + public boolean hasOpeningStreams(int nodeId) { + return streamControlManager.hasOpeningStreams(nodeId); + } + + @Override + public void lock(int nodeId) { + streamControlManager.lock(nodeId); + } + + @Override + public void unlock(int nodeId) { + streamControlManager.unlock(nodeId); + } +} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeCommittedEpoch.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeCommittedEpoch.java new file mode 100644 index 0000000000..af303048fc --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeCommittedEpoch.java @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.controller.stream; + +import org.apache.kafka.controller.automq.utils.AvroUtils; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; + +import java.io.IOException; + +import io.netty.buffer.ByteBuf; + +public class NodeCommittedEpoch { + public static final String NODE_COMMITED_EPOCH_KEY_PREFIX = "__a_r_c_nce/"; + private static final Schema SCHEMA0 = SchemaBuilder.record("NodeCommitedEpoch").fields() + .name("epoch").type().longType().noDefault() + .endRecord(); + private long epoch; + + public NodeCommittedEpoch(long epoch) { + this.epoch = epoch; + } + + public long getEpoch() { + return epoch; + } + + public void setEpoch(long epoch) { + this.epoch = epoch; + } + + @Override + public String toString() { + return "NodeCommitedEpoch{" + + "epoch=" + epoch + + '}'; + } + + public static ByteBuf encode(NodeCommittedEpoch nodeCommittedEpoch, int version) { + if (version != 0) { + throw new IllegalArgumentException("version must be 0"); + } + GenericRecord record = new GenericData.Record(SCHEMA0); + record.put("epoch", nodeCommittedEpoch.epoch); + + try { + return AvroUtils.encode(record, (short) 0); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + public static NodeCommittedEpoch decode(ByteBuf buf) { + try { + GenericRecord record = AvroUtils.decode(buf, version -> { + if (version != 0) { + throw new IllegalStateException("unsupported version: " + version); + } + return SCHEMA0; + }); + return new NodeCommittedEpoch((Long) record.get("epoch")); + } catch (IOException e) { + throw new RuntimeException(e); + } + + } +} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeControlManager.java index 6782cbcf77..b8a56cbd34 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeControlManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeControlManager.java @@ -1,16 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; +import org.apache.kafka.common.errors.s3.NodeLockedException; import org.apache.kafka.common.errors.s3.UnregisterNodeWithOpenStreamsException; import org.apache.kafka.common.message.AutomqGetNodesResponseData; import org.apache.kafka.common.message.AutomqRegisterNodeRequestData; @@ -24,6 +33,7 @@ import org.apache.kafka.server.common.ApiMessageAndVersion; import org.apache.kafka.timeline.SnapshotRegistry; import org.apache.kafka.timeline.TimelineHashMap; +import org.apache.kafka.timeline.TimelineHashSet; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -31,8 +41,10 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.stream.Collectors; public class NodeControlManager { @@ -40,12 +52,14 @@ public class NodeControlManager { private static final String KEY_PREFIX = "__automq_node/"; final TimelineHashMap nodeMetadataMap; + final TimelineHashSet lockedNodes; - private final NodeRuntimeInfoGetter nodeRuntimeInfoGetter; + private final NodeRuntimeInfoManager nodeRuntimeInfoManager; - public NodeControlManager(SnapshotRegistry registry, NodeRuntimeInfoGetter nodeRuntimeInfoGetter) { + public NodeControlManager(SnapshotRegistry registry, NodeRuntimeInfoManager nodeRuntimeInfoManager) { this.nodeMetadataMap = new TimelineHashMap<>(registry, 100); - this.nodeRuntimeInfoGetter = nodeRuntimeInfoGetter; + this.lockedNodes = new TimelineHashSet<>(registry, 100); + this.nodeRuntimeInfoManager = nodeRuntimeInfoManager; } public ControllerResult register(AutomqRegisterNodeRequest req) { @@ -102,17 +116,17 @@ public Collection getMetadata() { } public NodeState state(int nodeId) { - return nodeRuntimeInfoGetter.state(nodeId); + return nodeRuntimeInfoManager.state(nodeId); } /** * Note: It is costly to check if a node has opening streams, so it is recommended to use this method only when necessary. */ public boolean hasOpeningStreams(int nodeId) { - return nodeRuntimeInfoGetter.hasOpeningStreams(nodeId); + return nodeRuntimeInfoManager.hasOpeningStreams(nodeId); } - public void replay(KVRecord kvRecord) { + public synchronized void replay(KVRecord kvRecord) { for (KVRecord.KeyValue kv : kvRecord.keyValues()) { if (!(kv.key() != null && kv.key().startsWith(KEY_PREFIX))) { continue; @@ -121,12 +135,23 @@ public void replay(KVRecord kvRecord) { int nodeId = Integer.parseInt(kv.key().substring(KEY_PREFIX.length())); NodeMetadata nodeMetadata = NodeMetadataCodec.decode(kv.value()); nodeMetadataMap.put(nodeId, nodeMetadata); + if ("CLOSED".equals(nodeMetadata.getTags().getOrDefault("CIRCUIT_BREAKER", "CLOSED"))) { + nodeRuntimeInfoManager.unlock(nodeId); + lockedNodes.remove(nodeId); + } else { + nodeRuntimeInfoManager.lock(nodeId); + lockedNodes.add(nodeId); + } } catch (Throwable e) { LOGGER.error("[FATAL] replay NodeMetadata from KV fail", e); } } } + public synchronized Set lockedNodes() { + return new HashSet<>(lockedNodes); + } + ApiMessageAndVersion registerNodeRecord(int nodeId, NodeMetadata newNodeMetadata) { KVRecord kvRecord = new KVRecord().setKeyValues(List.of( new KVRecord.KeyValue() @@ -144,6 +169,7 @@ public void replay(RemoveKVRecord kvRecord) { try { int nodeId = Integer.parseInt(key.substring(KEY_PREFIX.length())); nodeMetadataMap.remove(nodeId); + nodeRuntimeInfoManager.unlock(nodeId); } catch (Throwable e) { LOGGER.error("[FATAL] replay NodeMetadata from KV fail", e); } @@ -154,6 +180,9 @@ public ApiMessageAndVersion unregisterNodeRecord(int nodeId) { if (hasOpeningStreams(nodeId)) { throw new UnregisterNodeWithOpenStreamsException(String.format("Node %d has opening streams", nodeId)); } + if (lockedNodes.contains(nodeId)) { + throw new NodeLockedException(String.format("Node %d is locked", nodeId)); + } RemoveKVRecord removeKVRecord = new RemoveKVRecord().setKeys(List.of(KEY_PREFIX + nodeId)); return new ApiMessageAndVersion(removeKVRecord, (short) 0); } diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadata.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadata.java index dc82780309..52e6c20dc5 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadata.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadata.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; @@ -17,6 +25,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import java.util.Map; +import java.util.Objects; @JsonIgnoreProperties(ignoreUnknown = true) public class NodeMetadata { @@ -85,6 +94,19 @@ public void setTags(Map tags) { this.tags = tags; } + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + NodeMetadata that = (NodeMetadata) o; + return nodeId == that.nodeId && nodeEpoch == that.nodeEpoch && Objects.equals(walConfig, that.walConfig) && Objects.equals(tags, that.tags); + } + + @Override + public int hashCode() { + return Objects.hash(nodeId, nodeEpoch, walConfig, tags); + } + @Override public String toString() { return "NodeMetadata{" + diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadataCodec.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadataCodec.java index d28cc0bd5e..d4a1966bd4 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadataCodec.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeMetadataCodec.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeRuntimeInfoGetter.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeRuntimeInfoGetter.java deleted file mode 100644 index 8251bc6343..0000000000 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeRuntimeInfoGetter.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package org.apache.kafka.controller.stream; - -public interface NodeRuntimeInfoGetter { - - NodeState state(int nodeId); - - boolean hasOpeningStreams(int nodeId); - -} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeRuntimeInfoManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeRuntimeInfoManager.java new file mode 100644 index 0000000000..e93a71cb57 --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeRuntimeInfoManager.java @@ -0,0 +1,32 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.controller.stream; + +public interface NodeRuntimeInfoManager { + + NodeState state(int nodeId); + + boolean hasOpeningStreams(int nodeId); + + void lock(int nodeId); + + void unlock(int nodeId); + +} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeState.java b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeState.java index 108b08d28e..48b5a3ae25 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/NodeState.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/NodeState.java @@ -1,16 +1,43 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; +import org.apache.kafka.controller.BrokerControlState; + public enum NodeState { - ACTIVE, FENCED, CONTROLLED_SHUTDOWN, UNKNOWN + /** + * The node is active and can handle requests. + */ + ACTIVE, + /** + * The node is shut down and cannot handle requests. + */ + FENCED, + /** + * The node is shutting down in a controlled manner. + * Note: In AutoMQ, this state is different from {@link BrokerControlState#CONTROLLED_SHUTDOWN}. In some cases, + * a node in {@link BrokerControlState#FENCED} state may still be shutting down in a controlled manner. + */ + CONTROLLED_SHUTDOWN, + /** + * The state of the node is unknown, possibly because it has not yet registered. + */ + UNKNOWN } diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java b/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java index 31f3ca60fc..52b5e0e39b 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * Use of this software is governed by the Business Source License - * included in the file BSL.md + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/RouterChannelEpoch.java b/metadata/src/main/java/org/apache/kafka/controller/stream/RouterChannelEpoch.java new file mode 100644 index 0000000000..1a13b3258d --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/RouterChannelEpoch.java @@ -0,0 +1,141 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.controller.stream; + +import org.apache.kafka.controller.automq.utils.AvroUtils; + +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; + +import java.io.IOException; +import java.util.Objects; + +import io.netty.buffer.ByteBuf; + +public class RouterChannelEpoch { + public static final String ROUTER_CHANNEL_EPOCH_KEY = "__a_r_c/metadata"; + private static final Schema SCHEMA0 = SchemaBuilder.record("RouterChannelEpoch").fields() + .name("committed").type().longType().noDefault() + .name("fenced").type().longType().noDefault() + .name("current").type().longType().noDefault() + .name("lastBumpUpTimestamp").type().longType().noDefault() + .endRecord(); + private long committed; + private long fenced; + private long current; + private long lastBumpUpTimestamp; + + public RouterChannelEpoch(long committed, long fenced, long current, long lastBumpUpTimestamp) { + this.committed = committed; + this.fenced = fenced; + this.current = current; + this.lastBumpUpTimestamp = lastBumpUpTimestamp; + } + + public long getCommitted() { + return committed; + } + + public void setCommitted(long committed) { + this.committed = committed; + } + + public long getFenced() { + return fenced; + } + + public void setFenced(long fenced) { + this.fenced = fenced; + } + + public long getCurrent() { + return current; + } + + public void setCurrent(long current) { + this.current = current; + } + + public long getLastBumpUpTimestamp() { + return lastBumpUpTimestamp; + } + + public void setLastBumpUpTimestamp(long lastBumpUpTimestamp) { + this.lastBumpUpTimestamp = lastBumpUpTimestamp; + } + + @Override + public String toString() { + return "RouterChannelEpoch{" + + "committed=" + committed + + ", fenced=" + fenced + + ", current=" + current + + ", lastBumpUpTimestamp=" + lastBumpUpTimestamp + + '}'; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + RouterChannelEpoch that = (RouterChannelEpoch) o; + return committed == that.committed && fenced == that.fenced && current == that.current && lastBumpUpTimestamp == that.lastBumpUpTimestamp; + } + + @Override + public int hashCode() { + return Objects.hash(committed, fenced, current, lastBumpUpTimestamp); + } + + public static ByteBuf encode(RouterChannelEpoch routerChannelEpoch, short version) { + if (version != 0) { + throw new IllegalArgumentException("version must be 0"); + } + GenericRecord record = new GenericData.Record(SCHEMA0); + record.put("committed", routerChannelEpoch.committed); + record.put("fenced", routerChannelEpoch.fenced); + record.put("current", routerChannelEpoch.current); + record.put("lastBumpUpTimestamp", routerChannelEpoch.lastBumpUpTimestamp); + + try { + return AvroUtils.encode(record, (short) 0); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + public static RouterChannelEpoch decode(ByteBuf buf) { + try { + GenericRecord record = AvroUtils.decode(buf, version -> { + if (version != 0) { + throw new IllegalStateException("unsupported version: " + version); + } + return SCHEMA0; + }); + return new RouterChannelEpoch( + (Long) record.get("committed"), + (Long) record.get("fenced"), + (Long) record.get("current"), + (Long) record.get("lastBumpUpTimestamp") + ); + } catch (IOException e) { + throw new RuntimeException(e); + } + } +} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/RouterChannelEpochControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/RouterChannelEpochControlManager.java new file mode 100644 index 0000000000..22b3b2ded0 --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/RouterChannelEpochControlManager.java @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.controller.stream; + +import org.apache.kafka.common.metadata.KVRecord; +import org.apache.kafka.common.metadata.RemoveKVRecord; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.controller.ControllerResult; +import org.apache.kafka.controller.QuorumController; +import org.apache.kafka.server.common.ApiMessageAndVersion; +import org.apache.kafka.timeline.SnapshotRegistry; +import org.apache.kafka.timeline.TimelineHashMap; +import org.apache.kafka.timeline.TimelineObject; + +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.LinkedList; +import java.util.List; +import java.util.OptionalLong; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import io.netty.buffer.Unpooled; + +public class RouterChannelEpochControlManager { + private static final Logger LOGGER = LoggerFactory.getLogger(RouterChannelEpochControlManager.class); + private static final long BUMP_EPOCH_INTERVAL = TimeUnit.SECONDS.toMillis(10); + private final TimelineObject routerChannelEpoch; + private final TimelineHashMap node2commitedEpoch; + private final ScheduledExecutorService scheduler = Threads.newSingleThreadScheduledExecutor("ROUTER_CHANNEL_EPOCH_MANAGER", true, LOGGER); + + private final QuorumController quorumController; + private final NodeControlManager nodeControlManager; + private final Time time; + + public RouterChannelEpochControlManager(SnapshotRegistry registry, QuorumController quorumController, NodeControlManager nodeControlManager, Time time) { + this.routerChannelEpoch = new TimelineObject<>(registry, new RouterChannelEpoch(-3L, -2L, 0, 0)); + this.node2commitedEpoch = new TimelineHashMap<>(registry, 100); + this.quorumController = quorumController; + this.nodeControlManager = nodeControlManager; + this.time = time; + scheduler.scheduleWithFixedDelay(this::run, 1, 1, TimeUnit.SECONDS); + } + + private void run() { + if (!quorumController.isActive()) { + return; + } + quorumController.appendWriteEvent("routerChannelEpochRun", OptionalLong.empty(), this::run0); + } + + private ControllerResult run0() { + List records = new LinkedList<>(); + + RouterChannelEpoch epoch = routerChannelEpoch.get(); + if (epoch == null) { + // initial the epoch if not exist + epoch = new RouterChannelEpoch(-3L, -2L, 0, time.milliseconds()); + records.add(new ApiMessageAndVersion( + kv(RouterChannelEpoch.ROUTER_CHANNEL_EPOCH_KEY, RouterChannelEpoch.encode(epoch, (short) 0).array()), (short) 0)); + } else { + // update the commitedEpoch, then RouterChannel can clean up commited data. + OptionalLong newCommitedEpoch = calCommitedEpoch(); + RouterChannelEpoch newEpoch = null; + if (time.milliseconds() - epoch.getLastBumpUpTimestamp() >= BUMP_EPOCH_INTERVAL) { + // bump to the next epoch. + newEpoch = new RouterChannelEpoch(newCommitedEpoch.orElse(epoch.getCommitted()), epoch.getFenced() + 1, epoch.getCurrent() + 1, time.milliseconds()); + } else if (newCommitedEpoch.isPresent() && newCommitedEpoch.getAsLong() > epoch.getCommitted()) { + newEpoch = new RouterChannelEpoch(newCommitedEpoch.getAsLong(), epoch.getFenced(), epoch.getCurrent(), epoch.getLastBumpUpTimestamp()); + } + if (newEpoch != null) { + records.add(new ApiMessageAndVersion( + kv(RouterChannelEpoch.ROUTER_CHANNEL_EPOCH_KEY, RouterChannelEpoch.encode(newEpoch, (short) 0).array()), (short) 0)); + } + } + + return ControllerResult.of(records, null); + } + + /** + * Calculate the newCommitedEpoch = min(nodeCommitedEpoch) + */ + private OptionalLong calCommitedEpoch() { + long newCommitedEpoch = Long.MAX_VALUE; + for (NodeMetadata nodeMetadata : nodeControlManager.getMetadata()) { + int nodeId = nodeMetadata.getNodeId(); + if (nodeControlManager.state(nodeId) == NodeState.FENCED && !nodeControlManager.hasOpeningStreams(nodeId)) { + // We can ignore the gracefully stopped node. + continue; + } + Long nodeCommitedEpoch = node2commitedEpoch.get(nodeId); + if (nodeCommitedEpoch == null) { + // The node + return OptionalLong.empty(); + } + newCommitedEpoch = Math.min(newCommitedEpoch, nodeCommitedEpoch); + } + if (newCommitedEpoch == Long.MAX_VALUE) { + return OptionalLong.empty(); + } + return OptionalLong.of(newCommitedEpoch); + } + + public void replay(KVRecord kvRecord) { + for (KVRecord.KeyValue kv : kvRecord.keyValues()) { + String key = kv.key(); + if (key == null) { + continue; + } + try { + if (key.startsWith(NodeCommittedEpoch.NODE_COMMITED_EPOCH_KEY_PREFIX)) { + int nodeId = Integer.parseInt(key.substring(NodeCommittedEpoch.NODE_COMMITED_EPOCH_KEY_PREFIX.length())); + NodeCommittedEpoch epoch = NodeCommittedEpoch.decode(Unpooled.wrappedBuffer(kv.value())); + node2commitedEpoch.put(nodeId, epoch.getEpoch()); + } else if (key.startsWith(RouterChannelEpoch.ROUTER_CHANNEL_EPOCH_KEY)) { + RouterChannelEpoch epoch = RouterChannelEpoch.decode(Unpooled.wrappedBuffer(kv.value())); + routerChannelEpoch.set(epoch); + } + } catch (Throwable e) { + LOGGER.error("[FATAL] replay router channel epoch {} fail", kv, e); + } + } + } + + public void replay(RemoveKVRecord kvRecord) { + for (String key : kvRecord.keys()) { + try { + if (key.startsWith(NodeCommittedEpoch.NODE_COMMITED_EPOCH_KEY_PREFIX)) { + int nodeId = Integer.parseInt(key.substring(NodeCommittedEpoch.NODE_COMMITED_EPOCH_KEY_PREFIX.length())); + node2commitedEpoch.remove(nodeId); + } + } catch (Throwable e) { + LOGGER.error("[FATAL] replay router channel epoch {} fail", key, e); + } + } + } + + static KVRecord kv(String key, byte[] value) { + return new KVRecord().setKeyValues(List.of(new KVRecord.KeyValue().setKey(key).setValue(value))); + } +} diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/S3ObjectControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/S3ObjectControlManager.java index 5d3804963c..2f71a1c50a 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/S3ObjectControlManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/S3ObjectControlManager.java @@ -48,6 +48,7 @@ import com.automq.stream.s3.objects.ObjectAttributes; import com.automq.stream.s3.objects.ObjectAttributes.Type; import com.automq.stream.s3.operator.AwsObjectStorage; +import com.automq.stream.s3.operator.LocalFileObjectStorage; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorage.ObjectPath; import com.automq.stream.utils.CollectionHelper; @@ -316,6 +317,26 @@ public ControllerResult markDestroyObjects(List objects, List replaceCommittedObject(long objectId, int attributes) { + S3Object object = this.objectsMetadata.get(objectId); + if (object == null) { + return ControllerResult.of(Collections.emptyList(), Errors.OBJECT_NOT_EXIST); + } + // verify the state + if (object.getS3ObjectState() != S3ObjectState.COMMITTED) { + return ControllerResult.of(Collections.emptyList(), Errors.OBJECT_NOT_COMMITED); + } + AutoMQVersion version = this.version.get(); + S3ObjectRecord record = new S3ObjectRecord() + .setObjectId(objectId) + .setObjectSize(object.getObjectSize()) + .setObjectState(S3ObjectState.COMMITTED.toByte()); + record.setTimestamp(object.getObjectSize()); + record.setAttributes(attributes); + return ControllerResult.of(List.of( + new ApiMessageAndVersion(record, version.objectRecordVersion())), Errors.NONE); + } + public void replay(AssignedS3ObjectIdRecord record) { nextAssignedObjectId.set(record.assignedS3ObjectId() + 1); } @@ -521,11 +542,14 @@ private void removeMarkDestroyedObject(long objectId) { class ObjectCleaner { CompletableFuture clean(List objects) { + List ignoredObjects = new LinkedList<>(); List deepDeleteCompositeObjects = new LinkedList<>(); List shallowDeleteObjects = new ArrayList<>(objects.size()); for (S3Object object : objects) { ObjectAttributes attributes = ObjectAttributes.from(object.getAttributes()); - if (attributes.deepDelete() && attributes.type() == Type.Composite) { + if (attributes.bucket() == LocalFileObjectStorage.BUCKET_ID) { + ignoredObjects.add(object); + } else if (attributes.deepDelete() && attributes.type() == Type.Composite) { deepDeleteCompositeObjects.add(object); } else { shallowDeleteObjects.add(object); @@ -537,6 +561,8 @@ CompletableFuture clean(List objects) { batchDelete(shallowDeleteObjects, this::shallowlyDelete, cfList); // Delete the composite object and it's linked objects batchDelete(deepDeleteCompositeObjects, this::deepDelete, cfList); + // Delete the local file objects + batchDelete(ignoredObjects, this::noopDelete, cfList); return CompletableFuture.allOf(cfList.toArray(new CompletableFuture[0])); } @@ -582,6 +608,11 @@ private CompletableFuture deepDelete(List s3Objects) { return allCf; } + private CompletableFuture noopDelete(List s3objects) { + List deletedObjectIds = s3objects.stream().map(S3Object::getObjectId).collect(Collectors.toList()); + return CompletableFuture.completedFuture(null).thenAccept(rst -> notifyS3ObjectDeleted(deletedObjectIds)); + } + private void notifyS3ObjectDeleted(List deletedObjectIds) { // notify the controller an objects deletion event to drive the removal of the objects ControllerRequestContext ctx = new ControllerRequestContext( diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamClient.java b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamClient.java index 58779c1a5b..9caf628238 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamClient.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamControlManager.java index a083429113..166b43e083 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamControlManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamControlManager.java @@ -18,6 +18,7 @@ package org.apache.kafka.controller.stream; import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData; import org.apache.kafka.common.message.CloseStreamsRequestData.CloseStreamRequest; import org.apache.kafka.common.message.CloseStreamsResponseData.CloseStreamResponse; import org.apache.kafka.common.message.CommitStreamObjectRequestData; @@ -56,6 +57,7 @@ import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.common.utils.ThreadUtils; import org.apache.kafka.controller.ClusterControlManager; +import org.apache.kafka.controller.ControllerRequestContext; import org.apache.kafka.controller.ControllerResult; import org.apache.kafka.controller.FeatureControlManager; import org.apache.kafka.controller.QuorumController; @@ -72,6 +74,7 @@ import org.apache.kafka.server.metrics.s3stream.S3StreamKafkaMetricsManager; import org.apache.kafka.timeline.SnapshotRegistry; import org.apache.kafka.timeline.TimelineHashMap; +import org.apache.kafka.timeline.TimelineHashSet; import org.apache.kafka.timeline.TimelineLong; import com.automq.stream.s3.ObjectReader; @@ -79,6 +82,9 @@ import com.automq.stream.s3.metadata.S3StreamConstant; import com.automq.stream.s3.metadata.StreamOffsetRange; import com.automq.stream.s3.metadata.StreamState; +import com.automq.stream.s3.objects.ObjectAttributes; +import com.automq.stream.s3.operator.LocalFileObjectStorage; +import com.google.common.base.Strings; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -122,6 +128,7 @@ public class StreamControlManager { private final TimelineHashMap streamsMetadata; private final TimelineHashMap nodesMetadata; + private final TimelineHashSet lockedNodes; private final TimelineHashMap stream2node; private final TimelineHashMap> node2streams; @@ -152,6 +159,7 @@ public StreamControlManager( this.nextAssignedStreamId = new TimelineLong(snapshotRegistry); this.streamsMetadata = new TimelineHashMap<>(snapshotRegistry, 100000); this.nodesMetadata = new TimelineHashMap<>(snapshotRegistry, 0); + this.lockedNodes = new TimelineHashSet<>(snapshotRegistry, 100); this.stream2node = new TimelineHashMap<>(snapshotRegistry, 100000); this.node2streams = new TimelineHashMap<>(snapshotRegistry, 100); @@ -291,6 +299,16 @@ public ControllerResult openStream(int nodeId, long nodeEpoc resp.setErrorCode(Errors.STREAM_NOT_CLOSED.code()); return ControllerResult.of(Collections.emptyList(), resp); } + int currentRangeOwner = streamMetadata.currentRangeOwner(); + if (nodeId != currentRangeOwner && lockedNodes.contains(currentRangeOwner)) { + // Forbidden other nodes to open the stream if the last range is owned by a locked node + resp.setErrorCode(Errors.NODE_LOCKED.code()); + log.warn("[OpenStream] the stream's last range is owned by a locked node {}. streamId={}, streamEpoch={}, requestEpoch={}, nodeId={}, nodeEpoch={}", + currentRangeOwner, streamId, streamMetadata.currentEpoch(), epoch, nodeId, nodeEpoch); + tryReassignPartitionBack(streamMetadata); + return ControllerResult.of(Collections.emptyList(), resp); + } + // now the request is valid, update the stream's epoch and create a new range for this node List records = new ArrayList<>(); int newRangeIndex = streamMetadata.currentRangeIndex() + 1; @@ -629,6 +647,11 @@ public ControllerResult commitStreamSetObject return ControllerResult.of(Collections.emptyList(), resp); } + if (data.compactedObjectIds().size() == 1 && data.objectId() == data.compactedObjectIds().get(0)) { + // replace the stream set object + return replace(data); + } + List streamRanges = data.objectStreamRanges(); List compactedObjectIds = data.compactedObjectIds(); List streamObjects = data.streamObjects(); @@ -649,7 +672,7 @@ public ControllerResult commitStreamSetObject List records = new ArrayList<>(commitResult.records()); long dataTs = committedTs; // mark destroy compacted object - if (compactedObjectIds != null && !compactedObjectIds.isEmpty()) { + if (!compactedObjectIds.isEmpty()) { ControllerResult destroyResult = this.s3ObjectControlManager.markDestroyObjects(compactedObjectIds); if (!destroyResult.response()) { log.error("[CommitStreamSetObject]: failed to mark destroy compacted objects. compactedObjects={}, streamSetObjectId={}, nodeId={}, nodeEpoch={}", @@ -660,10 +683,14 @@ public ControllerResult commitStreamSetObject records.addAll(destroyResult.records()); // update dataTs to the min compacted object's dataTs //noinspection OptionalGetWithoutIsPresent + NodeRuntimeMetadata nodeMetadata = this.nodesMetadata.get(nodeId); dataTs = compactedObjectIds.stream() - .map(id -> this.nodesMetadata.get(nodeId).streamSetObjects().get(id)) + .map(id -> nodeMetadata.streamSetObjects().get(id)) .map(S3StreamSetObject::dataTimeInMs) .min(Long::compareTo).get(); + if (orderId == -1L && !data.compactedObjectIds().isEmpty()) { + orderId = data.compactedObjectIds().stream().mapToLong(id -> nodeMetadata.streamSetObjects().get(id).orderId()).min().getAsLong(); + } } if (objectId != NOOP_OBJECT_ID) { // generate node's stream set object record @@ -671,18 +698,6 @@ public ControllerResult commitStreamSetObject if (version.isHugeClusterSupported()) { s3StreamSetObject = new S3StreamSetObject(objectId, nodeId, Bytes.EMPTY, orderId, dataTs); records.add(s3StreamSetObject.toRecord(version)); - // generate S3StreamEndOffsetsRecord to move stream endOffset - if (compactedObjectIds.isEmpty()) { - S3StreamEndOffsetsRecord record = new S3StreamEndOffsetsRecord().setEndOffsets( - S3StreamEndOffsetsCodec.encode( - Stream.concat( - streamRanges.stream().map(s -> new StreamEndOffset(s.streamId(), s.endOffset())), - streamObjects.stream().map(s -> new StreamEndOffset(s.streamId(), s.endOffset())) - ) - .collect(Collectors.toList())) - ); - records.add(new ApiMessageAndVersion(record, (short) 0)); - } } else { List indexes = streamRanges.stream() .map(range -> new StreamOffsetRange(range.streamId(), range.startOffset(), range.endOffset())) @@ -691,8 +706,20 @@ public ControllerResult commitStreamSetObject records.add(s3StreamSetObject.toRecord(version)); } } + if (compactedObjectIds.isEmpty() && version.isHugeClusterSupported()) { + // generate S3StreamEndOffsetsRecord to move stream endOffset + S3StreamEndOffsetsRecord record = new S3StreamEndOffsetsRecord().setEndOffsets( + S3StreamEndOffsetsCodec.encode( + Stream.concat( + streamRanges.stream().map(s -> new StreamEndOffset(s.streamId(), s.endOffset())), + streamObjects.stream().map(s -> new StreamEndOffset(s.streamId(), s.endOffset())) + ) + .collect(Collectors.toList())) + ); + records.add(new ApiMessageAndVersion(record, (short) 0)); + } // commit stream objects - if (streamObjects != null && !streamObjects.isEmpty()) { + if (!streamObjects.isEmpty()) { // commit objects ControllerResult ret = generateStreamObject(streamObjects, records, data, resp, committedTs); if (ret != null) { @@ -700,7 +727,7 @@ public ControllerResult commitStreamSetObject } } // generate compacted objects' remove record - if (compactedObjectIds != null && !compactedObjectIds.isEmpty()) { + if (!compactedObjectIds.isEmpty()) { compactedObjectIds.forEach(id -> records.add(new ApiMessageAndVersion(new RemoveStreamSetObjectRecord() .setNodeId(nodeId) .setObjectId(id), (short) 0))); @@ -715,6 +742,20 @@ public ControllerResult commitStreamSetObject return ControllerResult.atomicOf(records, resp); } + private ControllerResult replace(CommitStreamSetObjectRequestData data) { + CommitStreamSetObjectResponseData resp = new CommitStreamSetObjectResponseData(); + List records = new ArrayList<>(1); + long objectId = data.objectId(); + ControllerResult rst = s3ObjectControlManager.replaceCommittedObject(objectId, data.attributes()); + if (rst.response() == Errors.NONE) { + records.addAll(rst.records()); + return ControllerResult.of(records, resp); + } else { + resp.setErrorCode(rst.response().code()); + return ControllerResult.of(Collections.emptyList(), resp); + } + } + private ControllerResult generateStreamObject(List streamObjects, List records, CommitStreamSetObjectRequestData req, CommitStreamSetObjectResponseData resp, long committedTs) { @@ -802,6 +843,10 @@ public ControllerResult commitStreamObject(Commi CommitStreamObjectResponseData resp = new CommitStreamObjectResponseData(); long committedTs = System.currentTimeMillis(); + if (data.sourceObjectIds().size() == 1 && streamObjectId == data.sourceObjectIds().get(0)) { + return replace(data); + } + // verify node epoch Errors nodeEpochCheckResult = nodeEpochCheck(nodeId, nodeEpoch); if (nodeEpochCheckResult != Errors.NONE) { @@ -873,6 +918,20 @@ public ControllerResult commitStreamObject(Commi return ControllerResult.atomicOf(records, resp); } + private ControllerResult replace(CommitStreamObjectRequestData data) { + CommitStreamObjectResponseData resp = new CommitStreamObjectResponseData(); + List records = new ArrayList<>(1); + long objectId = data.objectId(); + ControllerResult rst = s3ObjectControlManager.replaceCommittedObject(objectId, data.attributes()); + if (rst.response() == Errors.NONE) { + records.addAll(rst.records()); + return ControllerResult.of(records, resp); + } else { + resp.setErrorCode(rst.response().code()); + return ControllerResult.of(Collections.emptyList(), resp); + } + } + private DescribeStreamsResponseData bulidDescribeStreamsResponseData( List streamRuntimeMetadataList) { List metadataList = streamRuntimeMetadataList.stream() @@ -1007,6 +1066,8 @@ public ControllerResult getOpeningStreams(GetOpen NodeRuntimeMetadata nodeRuntimeMetadata = this.nodesMetadata.get(nodeId); if (nodeRuntimeMetadata == null) { // create a new node metadata if absent + log.info("[GetOpeningStreams]: create new node metadata. nodeId={}, nodeEpoch={}, failoverMode={}", + nodeId, nodeEpoch, failoverMode); records.add(new ApiMessageAndVersion( new NodeWALMetadataRecord().setNodeId(nodeId).setNodeEpoch(nodeEpoch).setFailoverMode(failoverMode), (short) 0)); @@ -1016,12 +1077,15 @@ public ControllerResult getOpeningStreams(GetOpen if (nodeRuntimeMetadata != null && nodeEpoch < nodeRuntimeMetadata.getNodeEpoch()) { // node epoch has been expired resp.setErrorCode(Errors.NODE_EPOCH_EXPIRED.code()); - log.warn("[GetOpeningStreams]: expired node epoch. nodeId={}, nodeEpoch={}", nodeId, nodeEpoch); + log.warn("[GetOpeningStreams]: expired node epoch. nodeId={}, nodeEpoch={}, requestNodeEpoch={}", + nodeId, nodeRuntimeMetadata.getNodeEpoch(), nodeEpoch); return ControllerResult.of(Collections.emptyList(), resp); } if (nodeRuntimeMetadata != null) { // update node epoch + log.info("[GetOpeningStreams]: update node epoch. nodeId={}, oldNodeEpoch={}, newNodeEpoch={}, failoverMode={}", + nodeId, nodeRuntimeMetadata.getNodeEpoch(), nodeEpoch, failoverMode); records.add(new ApiMessageAndVersion( new NodeWALMetadataRecord().setNodeId(nodeId).setNodeEpoch(nodeEpoch).setFailoverMode(failoverMode), (short) 0)); @@ -1068,7 +1132,7 @@ private void doubleCheckOpeningStreams(List openingStreams, int public List getOpeningStreams(int nodeId) { List streamIdList = Optional.ofNullable(node2streams.get(nodeId)).map(l -> l.toList()).orElse(Collections.emptyList()); List streams = new ArrayList<>(streamIdList.size()); - for (Long streamId: streamIdList) { + for (Long streamId : streamIdList) { StreamRuntimeMetadata streamRuntimeMetadata = streamsMetadata.get(streamId); if (streamRuntimeMetadata == null) { continue; @@ -1234,7 +1298,12 @@ public ControllerResult cleanupScaleInNodes() { return; } List objects = new ArrayList<>(nodeRuntimeMetadata.streamSetObjects().values()); - if (objects.isEmpty()) { + boolean inMainStorageCircuitBreakerOpenStatus = objects.stream().anyMatch(sso -> { + return Optional.ofNullable(s3ObjectControlManager.getObject(sso.objectId())) + .map(o -> ObjectAttributes.from(o.getAttributes()).bucket() == LocalFileObjectStorage.BUCKET_ID) + .orElse(false); + }); + if (objects.isEmpty() || inMainStorageCircuitBreakerOpenStatus) { return; } CleanUpScaleInNodeContext ctx = new CleanUpScaleInNodeContext(nodeId, objects); @@ -1251,6 +1320,14 @@ public ControllerResult cleanupScaleInNodes() { return ControllerResult.of(records, null); } + public void lock(int nodeId) { + lockedNodes.add(nodeId); + } + + public void unlock(int nodeId) { + lockedNodes.remove(nodeId); + } + public void replay(AssignedStreamIdRecord record) { this.nextAssignedStreamId.set(record.assignedStreamId() + 1); } @@ -1550,4 +1627,34 @@ ControllerResult checkStreamSetObjectExpired(S3StreamSetObject object, LOGGER.info("clean up scaled-in node={} object={}", object.nodeId(), object.objectId()); return ControllerResult.of(records, true); } + + private void tryReassignPartitionBack(StreamRuntimeMetadata stream) { + ControllerRequestContext context = new ControllerRequestContext(null, null, OptionalLong.empty()); + AlterPartitionReassignmentsRequestData request = new AlterPartitionReassignmentsRequestData(); + String rawTopicId = stream.tags().get(StreamTags.Topic.KEY); + String rawPartitionIndex = stream.tags().get(StreamTags.Partition.KEY); + if (Strings.isNullOrEmpty(rawTopicId) || Strings.isNullOrEmpty(rawPartitionIndex)) { + return; + } + Uuid topicId = Uuid.fromString(rawTopicId); + int partitionIndex = StreamTags.Partition.decode(rawPartitionIndex); + int nodeId = stream.currentRangeOwner(); + quorumController.findTopicNames(context, List.of(topicId)).thenAccept(uuid2name -> { + String topicName = Optional.ofNullable(uuid2name.get(topicId)).filter(r -> !r.isError()).map(r -> r.result()).orElse(null); + if (topicName == null) { + return; + } + request.setTopics(List.of(new AlterPartitionReassignmentsRequestData.ReassignableTopic() + .setName(topicName) + .setPartitions(List.of( + new AlterPartitionReassignmentsRequestData.ReassignablePartition() + .setPartitionIndex(partitionIndex) + .setReplicas(List.of(nodeId)) + )))); + quorumController.alterPartitionReassignments(context, request) + .thenAccept(rst -> { + LOGGER.info("[REASSIGN_PARTITION_BACK_TO_LOCKED_NODE],req={},resp={}", request, rst); + }); + }); + } } diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamRuntimeMetadata.java b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamRuntimeMetadata.java index 95d1ea4fab..2ce4c6ae48 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/StreamRuntimeMetadata.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/StreamRuntimeMetadata.java @@ -95,6 +95,10 @@ public void currentRangeIndex(int currentRangeIndex) { * Return the owner (node id) of current range. */ public int currentRangeOwner() { + if (ranges.isEmpty()) { + // there is no range in a new stream + return -1; + } return ranges.get(currentRangeIndex.get()).nodeId(); } @@ -172,6 +176,11 @@ public List checkRemovableRanges() { RangeMetadata range = ranges.get(i); if (startOffset <= range.startOffset() && range.endOffset() <= endOffset) { removableRanges.add(range); + continue; + } + // remove middle ranges contains no message + if (range.startOffset() == range.endOffset()) { + removableRanges.add(range); } } return removableRanges; diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletion.java b/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletion.java index 6130088f13..50e9622603 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletion.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletion.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletionManager.java b/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletionManager.java index 66aa55b729..f9e30d6c8d 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletionManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/stream/TopicDeletionManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/main/java/org/apache/kafka/image/ClusterImage.java b/metadata/src/main/java/org/apache/kafka/image/ClusterImage.java index 35160c74fd..aa8bdfe6e9 100644 --- a/metadata/src/main/java/org/apache/kafka/image/ClusterImage.java +++ b/metadata/src/main/java/org/apache/kafka/image/ClusterImage.java @@ -17,17 +17,18 @@ package org.apache.kafka.image; +import org.apache.kafka.common.metadata.UpdateNextNodeIdRecord; import org.apache.kafka.image.node.ClusterImageNode; import org.apache.kafka.image.writer.ImageWriter; import org.apache.kafka.image.writer.ImageWriterOptions; import org.apache.kafka.metadata.BrokerRegistration; import org.apache.kafka.metadata.ControllerRegistration; +import org.apache.kafka.server.common.ApiMessageAndVersion; import java.util.Collections; import java.util.Map; import java.util.Objects; - /** * Represents the cluster in the metadata image. * @@ -100,6 +101,7 @@ public void write(ImageWriter writer, ImageWriterOptions options) { } } } + writer.write(new ApiMessageAndVersion(new UpdateNextNodeIdRecord().setNodeId(nextNodeId), (short) 0)); } @Override diff --git a/metadata/src/main/java/org/apache/kafka/image/DeltaList.java b/metadata/src/main/java/org/apache/kafka/image/DeltaList.java index b97aa39770..b5933a5805 100644 --- a/metadata/src/main/java/org/apache/kafka/image/DeltaList.java +++ b/metadata/src/main/java/org/apache/kafka/image/DeltaList.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image; diff --git a/metadata/src/main/java/org/apache/kafka/image/RegistryRef.java b/metadata/src/main/java/org/apache/kafka/image/RegistryRef.java index 45c4b43e5e..66300d8dea 100644 --- a/metadata/src/main/java/org/apache/kafka/image/RegistryRef.java +++ b/metadata/src/main/java/org/apache/kafka/image/RegistryRef.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image; diff --git a/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataDelta.java b/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataDelta.java index b879227d3c..adee0e27ab 100644 --- a/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataDelta.java +++ b/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataDelta.java @@ -108,6 +108,9 @@ public void replay(RemoveRangeRecord record) { public void replay(S3StreamObjectRecord record) { getOrCreateStreamMetadataDelta(record.streamId()).replay(record); + // https://github.com/AutoMQ/automq/issues/2333 + // try to fix the old stream end offset + updateStreamEndOffset(record.streamId(), record.endOffset()); } public void replay(RemoveS3StreamObjectRecord record) { @@ -132,6 +135,13 @@ public void replay(S3StreamEndOffsetsRecord record) { } } + public Set changedStreams() { + Set set = new HashSet<>(); + set.addAll(changedStreams.keySet()); + set.addAll(changedStreamEndOffsets.keySet()); + return set; + } + private void updateStreamEndOffset(long streamId, long newEndOffset) { changedStreamEndOffsets.compute(streamId, (id, offset) -> { if (offset == null) { diff --git a/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataImage.java b/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataImage.java index 6df1717478..928689b177 100644 --- a/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataImage.java +++ b/metadata/src/main/java/org/apache/kafka/image/S3StreamsMetadataImage.java @@ -55,6 +55,7 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; @@ -594,6 +595,7 @@ public int hashCode() { return Objects.hash(nextAssignedStreamId, streamMetadataList(), nodeMetadataList(), streamEndOffsets()); } + // caller use this value should be protected by registryRef lock public TimelineHashMap timelineNodeMetadata() { return nodeMetadataMap; } @@ -609,10 +611,15 @@ List nodeMetadataList() { }); } + // caller use this value should be protected by registryRef lock public TimelineHashMap timelineStreamMetadata() { return streamMetadataMap; } + public void inLockRun(Runnable runnable) { + registryRef.inLock(runnable); + } + List streamMetadataList() { if (registryRef == RegistryRef.NOOP) { return Collections.emptyList(); @@ -628,18 +635,34 @@ public long nextAssignedStreamId() { return nextAssignedStreamId; } + public OptionalLong streamEndOffset(long streamId) { + if (registryRef == RegistryRef.NOOP) { + return OptionalLong.empty(); + } + return registryRef.inLock(() -> { + Long endOffset = streamEndOffsets.get(streamId); + if (endOffset != null) { + return OptionalLong.of(endOffset); + } + // There is no record in a new stream + if (streamMetadataMap.containsKey(streamId)) { + return OptionalLong.of(0L); + } else { + return OptionalLong.empty(); + } + }); + } + + // caller use this value should be protected by registryRef lock TimelineHashMap> partition2streams() { return partition2streams; } + // caller use this value should be protected by registryRef lock TimelineHashMap stream2partition() { return stream2partition; } - RegistryRef registryRef() { - return registryRef; - } - // caller use this value should be protected by registryRef lock TimelineHashMap timelineStreamEndOffsets() { return streamEndOffsets; @@ -657,6 +680,10 @@ Map streamEndOffsets() { }); } + RegistryRef registryRef() { + return registryRef; + } + @Override public String toString() { return "S3StreamsMetadataImage{nextAssignedStreamId=" + nextAssignedStreamId + '}'; diff --git a/metadata/src/main/java/org/apache/kafka/image/StreamOffsetIndexMap.java b/metadata/src/main/java/org/apache/kafka/image/StreamOffsetIndexMap.java index fb23ceac25..2baa6ccb48 100644 --- a/metadata/src/main/java/org/apache/kafka/image/StreamOffsetIndexMap.java +++ b/metadata/src/main/java/org/apache/kafka/image/StreamOffsetIndexMap.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image; diff --git a/metadata/src/main/java/org/apache/kafka/image/TopicIdPartition.java b/metadata/src/main/java/org/apache/kafka/image/TopicIdPartition.java index cf65e5419c..be266ccc8b 100644 --- a/metadata/src/main/java/org/apache/kafka/image/TopicIdPartition.java +++ b/metadata/src/main/java/org/apache/kafka/image/TopicIdPartition.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image; diff --git a/metadata/src/main/java/org/apache/kafka/image/loader/MetadataListener.java b/metadata/src/main/java/org/apache/kafka/image/loader/MetadataListener.java index 77c9ab9c90..0234a84aa8 100644 --- a/metadata/src/main/java/org/apache/kafka/image/loader/MetadataListener.java +++ b/metadata/src/main/java/org/apache/kafka/image/loader/MetadataListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.loader; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/AutoMQImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/AutoMQImageNode.java index 59dc793e3d..a06e6339f0 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/AutoMQImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/AutoMQImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/KVImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/KVImageNode.java index baac166dac..79a322e6f9 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/KVImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/KVImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/NodeImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/NodeImageNode.java index e5b0bdbcc1..5e20586d3f 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/NodeImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/NodeImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/NodesImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/NodesImageNode.java index 1849a88b5a..e9eaa06811 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/NodesImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/NodesImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/ObjectsImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/ObjectsImageNode.java index 5bafdc0032..ad26300cb8 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/ObjectsImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/ObjectsImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamImageNode.java index 7111b93c02..eec8749217 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamsImageNode.java b/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamsImageNode.java index d9b0f1ceca..59dcbcd204 100644 --- a/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamsImageNode.java +++ b/metadata/src/main/java/org/apache/kafka/image/node/automq/StreamsImageNode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image.node.automq; diff --git a/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodec.java b/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodec.java index 93df594d61..d3319c9e7c 100644 --- a/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodec.java +++ b/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodec.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.metadata.stream; diff --git a/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamSetObject.java b/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamSetObject.java index 0a77052c1e..80a654ed29 100644 --- a/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamSetObject.java +++ b/metadata/src/main/java/org/apache/kafka/metadata/stream/S3StreamSetObject.java @@ -135,11 +135,11 @@ public static S3StreamSetObject of(S3StreamSetObjectRecord record) { record.ranges(), record.orderId(), record.dataTimeInMs()); } - public Integer nodeId() { + public int nodeId() { return nodeId; } - public Long objectId() { + public long objectId() { return objectId; } diff --git a/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamEndOffset.java b/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamEndOffset.java index 1516cf10bf..307bb952ba 100644 --- a/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamEndOffset.java +++ b/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamEndOffset.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.metadata.stream; diff --git a/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamTags.java b/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamTags.java index 6bdea0678e..6c993d2833 100644 --- a/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamTags.java +++ b/metadata/src/main/java/org/apache/kafka/metadata/stream/StreamTags.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.metadata.stream; diff --git a/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java index 9a8776f721..9c96de0b8b 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java @@ -24,6 +24,7 @@ import org.apache.kafka.controller.BrokerHeartbeatManager.BrokerHeartbeatStateIterator; import org.apache.kafka.controller.BrokerHeartbeatManager.BrokerHeartbeatStateList; import org.apache.kafka.controller.BrokerHeartbeatManager.UsableBrokerIterator; +import org.apache.kafka.controller.stream.NodeState; import org.apache.kafka.metadata.placement.UsableBroker; import org.junit.jupiter.api.Test; @@ -362,4 +363,43 @@ public void testTouchThrowsExceptionUnlessRegistered() { assertThrows(IllegalStateException.class, () -> manager.touch(4, false, 0)).getMessage()); } + + // AutoMQ inject start + @Test + public void testBrokerState() { + final long shutdownTimeoutNs = 10_000_000; // 10ms + // init + BrokerHeartbeatManager manager = newBrokerHeartbeatManager(); + manager.time().sleep(1000); + manager.register(0, true); + + // FENCED Broker + assertEquals(NodeState.FENCED, manager.brokerState(0, shutdownTimeoutNs)); + + // UNFENCED Broker + manager.touch(0, false, 100); + assertEquals(NodeState.ACTIVE, manager.brokerState(0, shutdownTimeoutNs)); + + // CONTROLLED_SHUTDOWN Broker + manager.maybeUpdateControlledShutdownOffset(0, 100); + assertEquals(NodeState.CONTROLLED_SHUTDOWN, manager.brokerState(0, shutdownTimeoutNs)); + + // SHUTDOWN_NOW Broker within shutdownTimeoutNs + manager.touch(0, true, 100); + manager.time().sleep(5); + assertEquals(NodeState.CONTROLLED_SHUTDOWN, manager.brokerState(0, shutdownTimeoutNs)); + + // SHUTDOWN_NOW Broker after shutdownTimeoutNs + manager.time().sleep(6); + assertEquals(NodeState.FENCED, manager.brokerState(0, shutdownTimeoutNs)); + + // UNFENCED Broker after SHUTDOWN + manager.touch(0, false, 100); + assertEquals(NodeState.ACTIVE, manager.brokerState(0, shutdownTimeoutNs)); + + // UNREGISTERED Broker + manager.remove(0); + assertEquals(NodeState.UNKNOWN, manager.brokerState(0, shutdownTimeoutNs)); + } + // AutoMQ inject end } diff --git a/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java index ddc7506b0d..4b9564a1f1 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java @@ -28,6 +28,7 @@ import org.apache.kafka.common.message.ControllerRegistrationRequestData; import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord; import org.apache.kafka.common.metadata.FenceBrokerRecord; +import org.apache.kafka.common.metadata.KVRecord; import org.apache.kafka.common.metadata.RegisterBrokerRecord; import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpoint; import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpointCollection; @@ -36,6 +37,7 @@ import org.apache.kafka.common.security.auth.SecurityProtocol; import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.common.utils.MockTime; +import org.apache.kafka.controller.stream.KVControlManager; import org.apache.kafka.image.writer.ImageWriterOptions; import org.apache.kafka.metadata.BrokerRegistration; import org.apache.kafka.metadata.BrokerRegistrationFencingChange; @@ -52,6 +54,7 @@ import org.apache.kafka.server.common.MetadataVersion; import org.apache.kafka.timeline.SnapshotRegistry; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.ParameterizedTest; @@ -60,12 +63,14 @@ import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Optional; +import java.util.Set; import java.util.stream.Stream; import static java.util.Arrays.asList; @@ -718,4 +723,57 @@ public void testReRegistrationAndBrokerEpoch(boolean newIncarnationId) { clusterControl.brokerRegistrations().get(1).epoch()); } } + + @Test + public void testReusableNodeIds() { + MockTime time = new MockTime(0, 0, 0); + SnapshotRegistry snapshotRegistry = new SnapshotRegistry(new LogContext()); + KVControlManager kvControl = new KVControlManager(snapshotRegistry, new LogContext()); + FeatureControlManager featureControl = new FeatureControlManager.Builder(). + setSnapshotRegistry(snapshotRegistry). + setQuorumFeatures(new QuorumFeatures(0, + QuorumFeatures.defaultFeatureMap(true), + Collections.singletonList(0))). + setMetadataVersion(MetadataVersion.IBP_3_9_IV0). + build(); + ClusterControlManager clusterControl = new ClusterControlManager.Builder(). + setTime(time). + setSnapshotRegistry(snapshotRegistry). + setSessionTimeoutNs(1000). + setFeatureControlManager(featureControl). + setBrokerUncleanShutdownHandler((brokerId, records) -> { }). + setQuorumVoters(new ArrayList<>()). + setKVControlManager(kvControl). + build(); + clusterControl.activate(); + Set nodeIds = clusterControl.getReusableNodeIds(); + Assertions.assertTrue(nodeIds.isEmpty()); + clusterControl.putReusableNodeIds(Set.of(1, 2, 3)).forEach(r -> { + kvControl.replay((KVRecord) r.message()); + }); + nodeIds = clusterControl.getReusableNodeIds(); + Assertions.assertEquals(Set.of(1, 2, 3), nodeIds); + + clusterControl.unRegisterBrokerRecords(4).forEach(r -> { + kvControl.replay((KVRecord) r.message()); + }); + nodeIds = clusterControl.getReusableNodeIds(); + Assertions.assertEquals(Set.of(1, 2, 3, 4), nodeIds); + + clusterControl.registerBrokerRecords(2).forEach(r -> { + kvControl.replay((KVRecord) r.message()); + }); + nodeIds = clusterControl.getReusableNodeIds(); + Assertions.assertEquals(Set.of(1, 3, 4), nodeIds); + + ControllerResult result = clusterControl.getNextNodeId(); + result.records().forEach(r -> { + if (r.message() instanceof KVRecord) { + kvControl.replay((KVRecord) r.message()); + } + }); + Set remainIds = new HashSet<>(Set.of(1, 3, 4)); + remainIds.remove(result.response()); + Assertions.assertEquals(remainIds, clusterControl.getReusableNodeIds()); + } } diff --git a/metadata/src/test/java/org/apache/kafka/controller/ElasticPartitionChangeBuilderTest.java b/metadata/src/test/java/org/apache/kafka/controller/ElasticPartitionChangeBuilderTest.java index 99e5fcfd16..6cb013eaaa 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/ElasticPartitionChangeBuilderTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/ElasticPartitionChangeBuilderTest.java @@ -28,10 +28,12 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import static org.junit.jupiter.api.Assertions.assertEquals; // TODO: add test for AutoMQ partition change +@Timeout(60) @Tag("S3Unit") public class ElasticPartitionChangeBuilderTest { @BeforeEach diff --git a/metadata/src/test/java/org/apache/kafka/controller/ElasticReplicationControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ElasticReplicationControlManagerTest.java index 2fac5d5ce5..31a93be33d 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/ElasticReplicationControlManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/ElasticReplicationControlManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller; diff --git a/metadata/src/test/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelectorTest.java b/metadata/src/test/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelectorTest.java index ee70a365ad..f31b0cc82e 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelectorTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/es/LoadAwarePartitionLeaderSelectorTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.es; diff --git a/metadata/src/test/java/org/apache/kafka/controller/stream/NodeControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/stream/NodeControlManagerTest.java index 7aadb9403a..531d450a8f 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/stream/NodeControlManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/stream/NodeControlManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; @@ -47,16 +55,16 @@ @Tag("S3Unit") public class NodeControlManagerTest { SnapshotRegistry registry; - NodeRuntimeInfoGetter nodeRuntimeInfoGetter; + NodeRuntimeInfoManager nodeRuntimeInfoManager; NodeControlManager nodeControlManager; @BeforeEach public void setup() { registry = new SnapshotRegistry(new LogContext()); - nodeRuntimeInfoGetter = mock(NodeRuntimeInfoGetter.class); + nodeRuntimeInfoManager = mock(NodeRuntimeInfoManager.class); - nodeControlManager = new NodeControlManager(registry, nodeRuntimeInfoGetter); + nodeControlManager = new NodeControlManager(registry, nodeRuntimeInfoManager); } @Test @@ -94,8 +102,8 @@ public void testRegister() { assertEquals(Errors.NONE.code(), rst.response().errorCode()); assertTrue(nodeControlManager.nodeMetadataMap.containsKey(0)); - when(nodeRuntimeInfoGetter.hasOpeningStreams(eq(0))).thenReturn(true); - when(nodeRuntimeInfoGetter.state(eq(0))).thenReturn(NodeState.FENCED); + when(nodeRuntimeInfoManager.hasOpeningStreams(eq(0))).thenReturn(true); + when(nodeRuntimeInfoManager.state(eq(0))).thenReturn(NodeState.FENCED); ControllerResult getRst = nodeControlManager.getMetadata( new AutomqGetNodesRequest(new AutomqGetNodesRequestData().setNodeIds(List.of(0, 1)), @@ -107,6 +115,7 @@ public void testRegister() { assertEquals(0, nodes.get(0).nodeId()); assertEquals(2L, nodes.get(0).nodeEpoch()); assertEquals("wal2", nodes.get(0).walConfig()); + assertEquals(NodeState.FENCED.name(), nodes.get(0).state()); } AutomqRegisterNodeRequestData.TagCollection tags(Map tags) { @@ -136,7 +145,7 @@ public void testUnregisterNodeWithOpenStreams() { assertTrue(nodeControlManager.nodeMetadataMap.containsKey(0)); // prepare: node has opening streams - when(nodeRuntimeInfoGetter.hasOpeningStreams(eq(0))).thenReturn(true); + when(nodeRuntimeInfoManager.hasOpeningStreams(eq(0))).thenReturn(true); // test: unregister node with open streams assertThrows(UnregisterNodeWithOpenStreamsException.class, () -> nodeControlManager.unregisterNodeRecord(0)); diff --git a/metadata/src/test/java/org/apache/kafka/controller/stream/OverloadCircuitBreakerTest.java b/metadata/src/test/java/org/apache/kafka/controller/stream/OverloadCircuitBreakerTest.java index fd913d897f..18c691f846 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/stream/OverloadCircuitBreakerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/stream/OverloadCircuitBreakerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * Use of this software is governed by the Business Source License - * included in the file BSL.md + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; @@ -16,10 +24,12 @@ import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; +@Timeout(60) @Tag("S3Unit") public class OverloadCircuitBreakerTest { diff --git a/metadata/src/test/java/org/apache/kafka/controller/stream/S3ObjectControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/stream/S3ObjectControlManagerTest.java index 7298ca0fd2..0351a78c17 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/stream/S3ObjectControlManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/stream/S3ObjectControlManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/test/java/org/apache/kafka/controller/stream/StreamControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/stream/StreamControlManagerTest.java index a13ac0a7c0..f908befb01 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/stream/StreamControlManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/stream/StreamControlManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; @@ -72,6 +80,7 @@ import com.automq.stream.s3.metadata.StreamState; import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.operator.MemoryObjectStorage; +import com.automq.stream.s3.operator.ObjectStorage; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; @@ -1523,12 +1532,12 @@ private void verifyFirstRange(StreamRuntimeMetadata streamMetadata, long expecte private Optional mockObjectReader( List ranges) throws ExecutionException, InterruptedException { MemoryObjectStorage objectStorage = new MemoryObjectStorage(); - ObjectWriter objectWriter = new ObjectWriter.DefaultObjectWriter(1, objectStorage, Integer.MAX_VALUE, Integer.MAX_VALUE); + ObjectWriter objectWriter = new ObjectWriter.DefaultObjectWriter(1, objectStorage, Integer.MAX_VALUE, Integer.MAX_VALUE, new ObjectStorage.WriteOptions()); ranges.forEach(range -> objectWriter.write( range.streamId(), List.of( - new StreamRecordBatch(range.streamId(), 0, range.startOffset(), (int) (range.endOffset() - range.startOffset()), Unpooled.buffer(1)) + StreamRecordBatch.of(range.streamId(), 0, range.startOffset(), (int) (range.endOffset() - range.startOffset()), Unpooled.buffer(1)) ) ) ); diff --git a/metadata/src/test/java/org/apache/kafka/controller/stream/StreamRuntimeMetadataTest.java b/metadata/src/test/java/org/apache/kafka/controller/stream/StreamRuntimeMetadataTest.java index 4976f26aa3..cb4ed3bf0d 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/stream/StreamRuntimeMetadataTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/stream/StreamRuntimeMetadataTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; @@ -21,12 +29,14 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.Collections; import java.util.List; import static org.junit.jupiter.api.Assertions.assertEquals; +@Timeout(60) @Tag("S3Unit") public class StreamRuntimeMetadataTest { private static final long STREAM_ID = 233L; diff --git a/metadata/src/test/java/org/apache/kafka/controller/stream/TopicDeletionManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/stream/TopicDeletionManagerTest.java index 7e71c9daa5..80761e7303 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/stream/TopicDeletionManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/stream/TopicDeletionManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.controller.stream; diff --git a/metadata/src/test/java/org/apache/kafka/image/DeltaListTest.java b/metadata/src/test/java/org/apache/kafka/image/DeltaListTest.java index 5658e46d1c..0c222aff79 100644 --- a/metadata/src/test/java/org/apache/kafka/image/DeltaListTest.java +++ b/metadata/src/test/java/org/apache/kafka/image/DeltaListTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image; diff --git a/metadata/src/test/java/org/apache/kafka/image/S3ObjectsImageTest.java b/metadata/src/test/java/org/apache/kafka/image/S3ObjectsImageTest.java index 2819362eb5..0ef9a4f8db 100644 --- a/metadata/src/test/java/org/apache/kafka/image/S3ObjectsImageTest.java +++ b/metadata/src/test/java/org/apache/kafka/image/S3ObjectsImageTest.java @@ -56,6 +56,14 @@ @Tag("S3Unit") public class S3ObjectsImageTest { + private static S3ObjectsImage getAndRetain(AtomicReference current, Object lock) { + synchronized (lock) { + S3ObjectsImage image = current.get(); + image.retain(); + return image; + } + } + static final S3ObjectsImage IMAGE1; static final List DELTA1_RECORDS; @@ -150,43 +158,54 @@ private void testToImageAndBack(S3ObjectsImage image) { public void testConcurrentRefRetainAndReleaseNotThrowException() throws InterruptedException { LogContext logContext = new LogContext("[S3ObjectsImageTest] "); Logger log = LoggerFactory.getLogger(getClass()); - + SnapshotRegistry registry = new SnapshotRegistry(logContext); + Object imageLock = new Object(); AtomicReference current = new AtomicReference<>(); TimelineHashMap map = new TimelineHashMap<>(registry, 10); - RegistryRef ref = new RegistryRef(registry, 0, new ArrayList<>()); - + List liveEpochs = new ArrayList<>(); + liveEpochs.add(0L); + RegistryRef ref = new RegistryRef(registry, 0, liveEpochs); + S3ObjectsImage start = new S3ObjectsImage(4L, map, ref); - current.set(start); - + synchronized (imageLock) { + current.set(start); + } + AtomicBoolean running = new AtomicBoolean(true); - + AtomicLong updateExceptionCounter = new AtomicLong(); AtomicLong accessExceptionCounter = new AtomicLong(); AtomicLong updateCounter = new AtomicLong(); AtomicLong accessCounter = new AtomicLong(); - + int threadCount = 9; // 1 update task + 8 access tasks CountDownLatch startLatch = new CountDownLatch(threadCount); CountDownLatch endLatch = new CountDownLatch(threadCount); - + ExecutorService es = Executors.newFixedThreadPool(threadCount); - + // Update task es.submit(() -> { startLatch.countDown(); try { startLatch.await(); while (running.get()) { - S3ObjectsImage image = current.get(); + S3ObjectsImage previous; + synchronized (imageLock) { + previous = current.get(); + } TimeUnit.MILLISECONDS.sleep(1); - current.set(new S3ObjectsImage(1, map, ref.next())); + S3ObjectsImage newImage = new S3ObjectsImage(1, map, ref.next()); + synchronized (imageLock) { + current.set(newImage); + } updateCounter.incrementAndGet(); TimeUnit.MILLISECONDS.sleep(1); - - if (image != current.get()) { + + if (previous != newImage) { try { - image.release(); + previous.release(); } catch (Throwable e) { updateExceptionCounter.incrementAndGet(); log.error("Exception in updateImageTask", e); @@ -199,17 +218,16 @@ public void testConcurrentRefRetainAndReleaseNotThrowException() throws Interrup endLatch.countDown(); } }); - - + + for (int i = 0; i < 8; i++) { es.submit(() -> { startLatch.countDown(); try { startLatch.await(); while (running.get()) { - S3ObjectsImage image = current.get(); + S3ObjectsImage image = getAndRetain(current, imageLock); try { - image.retain(); TimeUnit.MILLISECONDS.sleep(10); accessCounter.incrementAndGet(); } finally { @@ -228,22 +246,32 @@ public void testConcurrentRefRetainAndReleaseNotThrowException() throws Interrup } }); } - - startLatch.await(); + + startLatch.await(); TimeUnit.SECONDS.sleep(10); running.set(false); - - + + assertTrue(endLatch.await(5, TimeUnit.SECONDS), "Not all threads finished in time"); - + es.shutdownNow(); assertTrue(es.awaitTermination(5, TimeUnit.SECONDS), "ExecutorService did not terminate in time"); - + assertEquals(0, updateExceptionCounter.get(), "Exceptions in update task: " + updateExceptionCounter.get()); assertEquals(0, accessExceptionCounter.get(), "Exceptions in access tasks: " + accessExceptionCounter.get()); - + log.info("Update operations: {}", updateCounter.get()); log.info("Access operations: {}", accessCounter.get()); + + S3ObjectsImage finalImage; + synchronized (imageLock) { + finalImage = current.get(); + } + try { + finalImage.release(); + } catch (Throwable ignore) { + // best effort clean-up; errors here are irrelevant to the concurrency assertions + } } } diff --git a/metadata/src/test/java/org/apache/kafka/image/StreamOffsetIndexMapTest.java b/metadata/src/test/java/org/apache/kafka/image/StreamOffsetIndexMapTest.java index 0e3b21a9c3..cdc41701ee 100644 --- a/metadata/src/test/java/org/apache/kafka/image/StreamOffsetIndexMapTest.java +++ b/metadata/src/test/java/org/apache/kafka/image/StreamOffsetIndexMapTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.image; diff --git a/metadata/src/test/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodecTest.java b/metadata/src/test/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodecTest.java index 0c66f8a660..eaae5d8740 100644 --- a/metadata/src/test/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodecTest.java +++ b/metadata/src/test/java/org/apache/kafka/metadata/stream/S3StreamEndOffsetsCodecTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.metadata.stream; diff --git a/metadata/src/test/java/org/apache/kafka/metadata/stream/SortedStreamSetObjectsListTest.java b/metadata/src/test/java/org/apache/kafka/metadata/stream/SortedStreamSetObjectsListTest.java index 466773e845..a93f6f912f 100644 --- a/metadata/src/test/java/org/apache/kafka/metadata/stream/SortedStreamSetObjectsListTest.java +++ b/metadata/src/test/java/org/apache/kafka/metadata/stream/SortedStreamSetObjectsListTest.java @@ -19,6 +19,7 @@ import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import java.util.Collections; import java.util.List; @@ -26,6 +27,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; +@Timeout(60) @Tag("S3Unit") public class SortedStreamSetObjectsListTest { diff --git a/s3stream/README.md b/s3stream/README.md index 1418384767..af45f16969 100644 --- a/s3stream/README.md +++ b/s3stream/README.md @@ -1,6 +1,6 @@ ## S3Stream: A Shared Streaming Storage Library -S3Stream is a shared streaming storage library that provides a unified interface for reading and writing streaming data to cloud object storage services like Amazon S3, Google Cloud Storage, and Azure Blob Storage. EBS is utilized here for its low-latency capabilities. It is designed to be used as the storage layer for distributed systems like Apache Kafka, Apache RocketMQ, etc. It provides the following features: -* **High Reliability**: S3Stream leverages cloud storage services(EBS and S3) to achieve zero RPO, RTO in seconds and 99.999999999% durability. +S3Stream is a shared streaming storage library offering a unified interface for reading and writing streaming data to cloud object storage services such as Amazon S3, Google Cloud Storage, Azure Blob Storage, and any S3-compatible storage like MinIO. It is designed to be used as the storage layer for distributed streaming storage systems like Apache Kafka. It provides the following features: +* **High Reliability**: S3Stream leverages cloud storage services to achieve zero RPO, RTO in seconds and 99.999999999% durability. * **Cost Effective**: S3Stream is designed for optimal cost and efficiency on the cloud. It can cut Apache Kafka billing by 90% on the cloud. * **Unified Interface**: S3Stream provides a unified interface for reading and writing streaming data to cloud object storage services. * **High Performance**: S3Stream is optimized for high performance and low latency. It can handle high throughput and low latency workloads. @@ -43,19 +43,3 @@ public interface Stream { } ``` > Please refer to the [S3Stream API](src/main/java/com/automq/stream/api/Stream.java) for the newest API details. - -## S3Stream Architecture -![image](../docs/images/automq_s3stream_architecture.gif) - -In S3Stream's core architecture, data is initially written to the Write-Ahead Log (WAL) persistently, then it's uploaded to S3 storage in a near real-time fashion. To efficiently support two reading paradigms—Tailing Read and Catch-up Read—S3Stream incorporates a built-in Message Cache to expedite reading operations. -- **WAL Storage**: Opt for a storage medium with low latency; each WAL disk requires only a few GiB of space, with cloud storage like EBS typically being the choice. -- **S3 Storage**: Select the cloud provider's largest object storage service to offer high-throughput, cost-effective primary data storage solutions. -- **Message Cache**: Hot data and prefetched cold data are both stored in the cache to expedite reading. Simultaneously, they are efficiently evicted based on the consumer focus mechanism, thereby enhancing memory utilization efficiency. - -## Various WAL Storage Options -![image](../docs/images/automq_wal_architecture.gif) -S3Stream supports various WAL storage options, including EBS, Regional EBS, S3, and other cloud storage services. -- **EBS WAL**: EBS is the default choice for WAL storage, offering low latency and high durability. -- **Regional EBS WAL**: On Azure, GCP, and Alibaba Cloud, Regional EBS replicas span multiple AZs. -- **S3 WAL**: Utilizing S3 as a WAL eliminates the need for EBS, streamlining the architecture to be fully S3-based, thus simplifying operations and maintenance. If your current setup is limited to MinIO, this is an excellent option. -- **S3 Express WAL**: AWS provides S3 Express, a high-performance, low-latency object storage solution that is well-suited as a storage choice for the S3Stream WAL. \ No newline at end of file diff --git a/s3stream/pom.xml b/s3stream/pom.xml index 2e4ba91d1c..dbbe6d73bf 100644 --- a/s3stream/pom.xml +++ b/s3stream/pom.xml @@ -1,13 +1,6 @@ HdrHistogram 2.2.2 + + com.ibm.async + asyncutil + 0.1.0 + diff --git a/s3stream/src/main/java/com/automq/stream/ByteBufSeqAlloc.java b/s3stream/src/main/java/com/automq/stream/ByteBufSeqAlloc.java index 62c649b2a7..0d09550b9c 100644 --- a/s3stream/src/main/java/com/automq/stream/ByteBufSeqAlloc.java +++ b/s3stream/src/main/java/com/automq/stream/ByteBufSeqAlloc.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/main/java/com/automq/stream/Context.java b/s3stream/src/main/java/com/automq/stream/Context.java new file mode 100644 index 0000000000..8596dfcb93 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/Context.java @@ -0,0 +1,61 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream; + +import com.automq.stream.api.KVClient; +import com.automq.stream.s3.ConfirmWAL; +import com.automq.stream.s3.cache.SnapshotReadCache; + +public class Context { + private SnapshotReadCache snapshotReadCache; + private ConfirmWAL confirmWAL; + private KVClient kvClient; + + public static final Context INSTANCE = new Context(); + + public static Context instance() { + return INSTANCE; + } + + public KVClient kvClient() { + return kvClient; + } + + public void kvClient(KVClient kvClient) { + this.kvClient = kvClient; + } + + public void snapshotReadCache(SnapshotReadCache snapshotReadCache) { + this.snapshotReadCache = snapshotReadCache; + } + + public SnapshotReadCache snapshotReadCache() { + return snapshotReadCache; + } + + public void confirmWAL(ConfirmWAL confirmWAL) { + this.confirmWAL = confirmWAL; + } + + public ConfirmWAL confirmWAL() { + return confirmWAL; + } + +} diff --git a/s3stream/src/main/java/com/automq/stream/DefaultAppendResult.java b/s3stream/src/main/java/com/automq/stream/DefaultAppendResult.java index 5857ef90d7..8f49893e68 100644 --- a/s3stream/src/main/java/com/automq/stream/DefaultAppendResult.java +++ b/s3stream/src/main/java/com/automq/stream/DefaultAppendResult.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/main/java/com/automq/stream/DefaultRecordBatch.java b/s3stream/src/main/java/com/automq/stream/DefaultRecordBatch.java index 792a067e71..5c540118e3 100644 --- a/s3stream/src/main/java/com/automq/stream/DefaultRecordBatch.java +++ b/s3stream/src/main/java/com/automq/stream/DefaultRecordBatch.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java b/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java index 6ec1f4dfd4..daa4beed3a 100644 --- a/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java +++ b/s3stream/src/main/java/com/automq/stream/FixedSizeByteBufPool.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/main/java/com/automq/stream/RecordBatchWithContextWrapper.java b/s3stream/src/main/java/com/automq/stream/RecordBatchWithContextWrapper.java index 14562bea99..9186219ba5 100644 --- a/s3stream/src/main/java/com/automq/stream/RecordBatchWithContextWrapper.java +++ b/s3stream/src/main/java/com/automq/stream/RecordBatchWithContextWrapper.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/main/java/com/automq/stream/Version.java b/s3stream/src/main/java/com/automq/stream/Version.java index 4e5cae1f1b..49cf5680ce 100644 --- a/s3stream/src/main/java/com/automq/stream/Version.java +++ b/s3stream/src/main/java/com/automq/stream/Version.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/main/java/com/automq/stream/api/AppendResult.java b/s3stream/src/main/java/com/automq/stream/api/AppendResult.java index 711aae9dc5..f1703305e9 100644 --- a/s3stream/src/main/java/com/automq/stream/api/AppendResult.java +++ b/s3stream/src/main/java/com/automq/stream/api/AppendResult.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/Client.java b/s3stream/src/main/java/com/automq/stream/api/Client.java index df1c9d71c3..93dc62828a 100644 --- a/s3stream/src/main/java/com/automq/stream/api/Client.java +++ b/s3stream/src/main/java/com/automq/stream/api/Client.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/CreateStreamOptions.java b/s3stream/src/main/java/com/automq/stream/api/CreateStreamOptions.java index bfc5be5711..70cc75311d 100644 --- a/s3stream/src/main/java/com/automq/stream/api/CreateStreamOptions.java +++ b/s3stream/src/main/java/com/automq/stream/api/CreateStreamOptions.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/FetchResult.java b/s3stream/src/main/java/com/automq/stream/api/FetchResult.java index 7548712bd2..e9a130c3d4 100644 --- a/s3stream/src/main/java/com/automq/stream/api/FetchResult.java +++ b/s3stream/src/main/java/com/automq/stream/api/FetchResult.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/KVClient.java b/s3stream/src/main/java/com/automq/stream/api/KVClient.java index b4a5c0d998..bf6f5b1441 100644 --- a/s3stream/src/main/java/com/automq/stream/api/KVClient.java +++ b/s3stream/src/main/java/com/automq/stream/api/KVClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; @@ -32,7 +40,7 @@ public interface KVClient { * Put key value, overwrite if key exist, return current key value after putting. * * @param keyValue {@link KeyValue} k-v pair - * @return async put result. {@link KeyValue} current value after putting. + * @return async put result. {@link Value} current value after putting. */ CompletableFuture putKV(KeyValue keyValue); @@ -40,7 +48,7 @@ public interface KVClient { * Get value by key. * * @param key key. - * @return async get result. {@link KeyValue} k-v pair, null if key not exist. + * @return async get result. {@link Value} retrieved value, null if key not exist. */ CompletableFuture getKV(Key key); diff --git a/s3stream/src/main/java/com/automq/stream/api/KeyValue.java b/s3stream/src/main/java/com/automq/stream/api/KeyValue.java index 3a91438976..7722120c48 100644 --- a/s3stream/src/main/java/com/automq/stream/api/KeyValue.java +++ b/s3stream/src/main/java/com/automq/stream/api/KeyValue.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/LinkRecordDecoder.java b/s3stream/src/main/java/com/automq/stream/api/LinkRecordDecoder.java new file mode 100644 index 0000000000..2add81c427 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/api/LinkRecordDecoder.java @@ -0,0 +1,51 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.api; + +import com.automq.stream.s3.model.StreamRecordBatch; + +import java.util.concurrent.CompletableFuture; + +import io.netty.buffer.ByteBuf; + +public interface LinkRecordDecoder { + LinkRecordDecoder NOOP = new Noop(); + + /** + * Get the decoded record size + */ + int decodedSize(ByteBuf linkRecordBuf); + + CompletableFuture decode(StreamRecordBatch src); + + + class Noop implements LinkRecordDecoder { + + @Override + public int decodedSize(ByteBuf linkRecordBuf) { + throw new UnsupportedOperationException(); + } + + @Override + public CompletableFuture decode(StreamRecordBatch src) { + return CompletableFuture.failedFuture(new UnsupportedOperationException()); + } + } +} diff --git a/s3stream/src/main/java/com/automq/stream/api/OpenStreamOptions.java b/s3stream/src/main/java/com/automq/stream/api/OpenStreamOptions.java index b8524aaf71..2a45202af3 100644 --- a/s3stream/src/main/java/com/automq/stream/api/OpenStreamOptions.java +++ b/s3stream/src/main/java/com/automq/stream/api/OpenStreamOptions.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; @@ -17,8 +25,9 @@ import java.util.Map; public class OpenStreamOptions { - private WriteMode writeMode = WriteMode.SINGLE; - private ReadMode readMode = ReadMode.MULTIPLE; + public static final OpenStreamOptions DEFAULT = new OpenStreamOptions(); + + private ReadWriteMode readWriteMode = ReadWriteMode.READ_WRITE; private long epoch; private final Map tags = new HashMap<>(); @@ -29,12 +38,8 @@ public static Builder builder() { return new Builder(); } - public WriteMode writeMode() { - return writeMode; - } - - public ReadMode readMode() { - return readMode; + public ReadWriteMode readWriteMode() { + return readWriteMode; } public long epoch() { @@ -45,26 +50,12 @@ public Map tags() { return tags; } - public enum WriteMode { - SINGLE(0), MULTIPLE(1); + public enum ReadWriteMode { + READ_WRITE(0), SNAPSHOT_READ(1); final int code; - WriteMode(int code) { - this.code = code; - } - - public int getCode() { - return code; - } - } - - public enum ReadMode { - SINGLE(0), MULTIPLE(1); - - final int code; - - ReadMode(int code) { + ReadWriteMode(int code) { this.code = code; } @@ -76,15 +67,10 @@ public int getCode() { public static class Builder { private final OpenStreamOptions options = new OpenStreamOptions(); - public Builder writeMode(WriteMode writeMode) { - Arguments.isNotNull(writeMode, "WriteMode should be set with SINGLE or MULTIPLE"); - options.writeMode = writeMode; - return this; - } + public Builder readWriteMode(ReadWriteMode readWriteMode) { + Arguments.isNotNull(readWriteMode, "readWriteMode should be set with READ_WRITE or SNAPSHOT_READ"); - public Builder readMode(ReadMode readMode) { - Arguments.isNotNull(readMode, "ReadMode should be set with SINGLE or MULTIPLE"); - options.readMode = readMode; + options.readWriteMode = readWriteMode; return this; } diff --git a/s3stream/src/main/java/com/automq/stream/api/ReadOptions.java b/s3stream/src/main/java/com/automq/stream/api/ReadOptions.java index 8aa76fdff2..d4a135c591 100644 --- a/s3stream/src/main/java/com/automq/stream/api/ReadOptions.java +++ b/s3stream/src/main/java/com/automq/stream/api/ReadOptions.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; @@ -19,6 +27,7 @@ public class ReadOptions { private boolean fastRead; private boolean pooledBuf; private boolean prioritizedRead; + private boolean snapshotRead; public static Builder builder() { return new Builder(); @@ -36,6 +45,15 @@ public boolean prioritizedRead() { return prioritizedRead; } + public boolean snapshotRead() { + return snapshotRead; + } + + public ReadOptions snapshotRead(boolean snapshotRead) { + this.snapshotRead = snapshotRead; + return this; + } + public static class Builder { private final ReadOptions options = new ReadOptions(); @@ -60,6 +78,11 @@ public Builder prioritizedRead(boolean prioritizedRead) { return this; } + public Builder snapshotRead(boolean snapshotRead) { + options.snapshotRead = snapshotRead; + return this; + } + public ReadOptions build() { return options; } diff --git a/s3stream/src/main/java/com/automq/stream/api/RecordBatch.java b/s3stream/src/main/java/com/automq/stream/api/RecordBatch.java index d8da445b24..4e34779882 100644 --- a/s3stream/src/main/java/com/automq/stream/api/RecordBatch.java +++ b/s3stream/src/main/java/com/automq/stream/api/RecordBatch.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/RecordBatchWithContext.java b/s3stream/src/main/java/com/automq/stream/api/RecordBatchWithContext.java index 1e9b32d075..5332abd327 100644 --- a/s3stream/src/main/java/com/automq/stream/api/RecordBatchWithContext.java +++ b/s3stream/src/main/java/com/automq/stream/api/RecordBatchWithContext.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/Stream.java b/s3stream/src/main/java/com/automq/stream/api/Stream.java index e440e4ffd1..5e25faeb46 100644 --- a/s3stream/src/main/java/com/automq/stream/api/Stream.java +++ b/s3stream/src/main/java/com/automq/stream/api/Stream.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; @@ -42,6 +50,11 @@ public interface Stream { */ long confirmOffset(); + /** + * Set confirm offset. Only support in snapshotRead mode + */ + void confirmOffset(long offset); + /** * Get stream next append record offset. */ @@ -95,4 +108,9 @@ default CompletableFuture fetch(long startOffset, long endOffset, i * Destroy stream. */ CompletableFuture destroy(); + + /** + * Get last append future + */ + CompletableFuture lastAppendFuture(); } diff --git a/s3stream/src/main/java/com/automq/stream/api/StreamClient.java b/s3stream/src/main/java/com/automq/stream/api/StreamClient.java index b539afa17f..8ac8672c83 100644 --- a/s3stream/src/main/java/com/automq/stream/api/StreamClient.java +++ b/s3stream/src/main/java/com/automq/stream/api/StreamClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api; diff --git a/s3stream/src/main/java/com/automq/stream/api/exceptions/ErrorCode.java b/s3stream/src/main/java/com/automq/stream/api/exceptions/ErrorCode.java index 28f7bc5679..06406ea28b 100644 --- a/s3stream/src/main/java/com/automq/stream/api/exceptions/ErrorCode.java +++ b/s3stream/src/main/java/com/automq/stream/api/exceptions/ErrorCode.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api.exceptions; diff --git a/s3stream/src/main/java/com/automq/stream/api/exceptions/FastReadFailFastException.java b/s3stream/src/main/java/com/automq/stream/api/exceptions/FastReadFailFastException.java index 3747e3e6a5..504af24fa7 100644 --- a/s3stream/src/main/java/com/automq/stream/api/exceptions/FastReadFailFastException.java +++ b/s3stream/src/main/java/com/automq/stream/api/exceptions/FastReadFailFastException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api.exceptions; diff --git a/s3stream/src/main/java/com/automq/stream/api/exceptions/StreamClientException.java b/s3stream/src/main/java/com/automq/stream/api/exceptions/StreamClientException.java index c7ac9678c0..823a2e2216 100644 --- a/s3stream/src/main/java/com/automq/stream/api/exceptions/StreamClientException.java +++ b/s3stream/src/main/java/com/automq/stream/api/exceptions/StreamClientException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.api.exceptions; diff --git a/s3stream/src/main/java/com/automq/stream/s3/ByteBufAlloc.java b/s3stream/src/main/java/com/automq/stream/s3/ByteBufAlloc.java index 3d4394e599..eb669f87f4 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/ByteBufAlloc.java +++ b/s3stream/src/main/java/com/automq/stream/s3/ByteBufAlloc.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -52,6 +60,7 @@ public class ByteBufAlloc { public static final int BLOCK_CACHE = 11; public static final int S3_WAL = 12; public static final int POOLED_MEMORY_RECORDS = 13; + public static final int SNAPSHOT_READ_CACHE = 14; // the MAX_TYPE_NUMBER may change when new type added. public static final int MAX_TYPE_NUMBER = 20; @@ -96,6 +105,7 @@ public class ByteBufAlloc { registerAllocType(BLOCK_CACHE, "block_cache"); registerAllocType(S3_WAL, "s3_wal"); registerAllocType(POOLED_MEMORY_RECORDS, "pooled_memory_records"); + registerAllocType(SNAPSHOT_READ_CACHE, "snapshot_read_cache"); } @@ -137,7 +147,7 @@ public static ByteBuf byteBuffer(int initCapacity, int type) { if (MEMORY_USAGE_DETECT) { LongAdder counter; - if (type > MAX_TYPE_NUMBER) { + if (type >= MAX_TYPE_NUMBER || type < 0) { counter = UNKNOWN_USAGE_STATS; } else { counter = USAGE_STATS[type]; diff --git a/s3stream/src/main/java/com/automq/stream/s3/ByteBufAllocPolicy.java b/s3stream/src/main/java/com/automq/stream/s3/ByteBufAllocPolicy.java index 74faaba5bf..db662aba3f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/ByteBufAllocPolicy.java +++ b/s3stream/src/main/java/com/automq/stream/s3/ByteBufAllocPolicy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/CompositeObject.java b/s3stream/src/main/java/com/automq/stream/s3/CompositeObject.java index 10a1580809..8574034338 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/CompositeObject.java +++ b/s3stream/src/main/java/com/automq/stream/s3/CompositeObject.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectReader.java b/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectReader.java index a04e7c6ef3..5e1c9ffce5 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectReader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectReader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -25,6 +33,7 @@ import java.util.List; import java.util.Objects; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import io.netty.buffer.ByteBuf; @@ -44,6 +53,7 @@ public class CompositeObjectReader implements ObjectReader { private CompletableFuture basicObjectInfoCf; private CompletableFuture sizeCf; private final AtomicInteger refCount = new AtomicInteger(1); + private final AtomicBoolean isShutdown = new AtomicBoolean(false); public CompositeObjectReader(S3ObjectMetadata objectMetadata, RangeReader rangeReader) { this.objectMetadata = objectMetadata; @@ -62,6 +72,9 @@ public String objectKey() { @Override public synchronized CompletableFuture basicObjectInfo() { + if (isShutdown.get()) { + return CompletableFuture.failedFuture(new IllegalStateException("ObjectReader is already shutdown")); + } if (basicObjectInfoCf == null) { this.basicObjectInfoCf = new CompletableFuture<>(); this.basicObjectInfoCf.exceptionally(ex -> { @@ -104,6 +117,9 @@ public synchronized CompletableFuture size() { } public synchronized void close0() { + if (!isShutdown.compareAndSet(false, true)) { + return; + } if (basicObjectInfoCf != null) { basicObjectInfoCf.thenAccept(BasicObjectInfo::close); } @@ -111,17 +127,14 @@ public synchronized void close0() { @Override public boolean equals(Object o) { - if (this == o) - return true; - if (o == null || getClass() != o.getClass()) - return false; - CompositeObjectReader reader = (CompositeObjectReader) o; - return Objects.equals(objectMetadata.objectId(), reader.objectMetadata.objectId()); + // NOTE: DO NOT OVERRIDE THIS + return super.equals(o); } @Override public int hashCode() { - return Objects.hashCode(objectMetadata.objectId()); + // NOTE: DO NOT OVERRIDE THIS + return super.hashCode(); } private void asyncGetBasicObjectInfo(CompletableFuture basicObjectInfoCf) { diff --git a/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectWriter.java b/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectWriter.java index 04ee68781e..2b9756a7ed 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectWriter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/CompositeObjectWriter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/Config.java b/s3stream/src/main/java/com/automq/stream/s3/Config.java index 0fbcfeea8d..d1c7dd0a03 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/Config.java +++ b/s3stream/src/main/java/com/automq/stream/s3/Config.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -25,6 +33,8 @@ public class Config { private String walConfig = "0@file:///tmp/s3stream_wal"; private long walCacheSize = 200 * 1024 * 1024; private long walUploadThreshold = 100 * 1024 * 1024; + // -1L means don't upload by time + private long walUploadIntervalMs = -1L; private int streamSplitSize = 16777216; private int objectBlockSize = 1048576; private int objectPartSize = 16777216; @@ -35,20 +45,21 @@ public class Config { private int controllerRequestRetryMaxCount = Integer.MAX_VALUE; private long controllerRequestRetryBaseDelayMs = 500; private long nodeEpoch = 0L; - private int streamSetObjectCompactionInterval = 10; + private int streamSetObjectCompactionInterval = 5; private long streamSetObjectCompactionCacheSize = 200 * 1024 * 1024; private int streamSetObjectCompactionUploadConcurrency = 8; - private long streamSetObjectCompactionStreamSplitSize = 16 * 1024 * 1024; + private long streamSetObjectCompactionStreamSplitSize = 8 * 1024 * 1024; private int streamSetObjectCompactionForceSplitPeriod = 120; private int streamSetObjectCompactionMaxObjectNum = 500; - private int maxStreamNumPerStreamSetObject = 100000; + private int maxStreamNumPerStreamSetObject = 20000; private int maxStreamObjectNumPerCommit = 10000; private boolean mockEnable = false; - // 100MB/s - private long networkBaselineBandwidth = 100 * 1024 * 1024; + // 1GBps/s + private long networkBaselineBandwidth = 1024 * 1024 * 1024; private int refillPeriodMs = 10; private long objectRetentionTimeInSecond = 10 * 60; // 10min private boolean failoverEnable = false; + private boolean snapshotReadEnable = false; private Supplier version = () -> { throw new UnsupportedOperationException(); }; @@ -73,6 +84,10 @@ public long walUploadThreshold() { return walUploadThreshold; } + public long walUploadIntervalMs() { + return walUploadIntervalMs; + } + public int streamSplitSize() { return streamSplitSize; } @@ -182,6 +197,11 @@ public Config walUploadThreshold(long s3WALObjectSize) { return this; } + public Config walUploadIntervalMs(long s3WALUploadIntervalMs) { + this.walUploadIntervalMs = s3WALUploadIntervalMs; + return this; + } + public Config streamSplitSize(int s3StreamSplitSize) { this.streamSplitSize = s3StreamSplitSize; return this; @@ -305,6 +325,15 @@ public boolean failoverEnable() { return failoverEnable; } + public Config snapshotReadEnable(boolean snapshotReadEnable) { + this.snapshotReadEnable = snapshotReadEnable; + return this; + } + + public boolean snapshotReadEnable() { + return snapshotReadEnable; + } + public Config version(Supplier version) { this.version = version; return this; diff --git a/s3stream/src/main/java/com/automq/stream/s3/ConfigValidator.java b/s3stream/src/main/java/com/automq/stream/s3/ConfigValidator.java new file mode 100644 index 0000000000..8bf388c9db --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/ConfigValidator.java @@ -0,0 +1,37 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3; + +import io.netty.util.internal.PlatformDependent; + +public class ConfigValidator { + + public static void validate(Config config) { + long memoryLimit = ByteBufAlloc.getPolicy().isDirect() ? PlatformDependent.maxDirectMemory() : Runtime.getRuntime().maxMemory(); + long memoryRequired = config.blockCacheSize() + config.walCacheSize(); + if (memoryRequired > memoryLimit) { + throw new IllegalArgumentException(String.format("blockCacheSize + walCacheSize size %s exceeds %s limit of %s", memoryRequired, ByteBufAlloc.getPolicy(), memoryLimit)); + } + if (config.walUploadThreshold() > config.walCacheSize()) { + throw new IllegalArgumentException(String.format("walUploadThreshold %s exceeds walCacheSize %s", config.walUploadThreshold(), config.walCacheSize())); + } + } + +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/ConfirmWAL.java b/s3stream/src/main/java/com/automq/stream/s3/ConfirmWAL.java new file mode 100644 index 0000000000..3b5b5a9100 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/ConfirmWAL.java @@ -0,0 +1,82 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3; + +import com.automq.stream.s3.S3Storage.LazyCommit; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.WriteAheadLog; + +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.function.Function; + +public class ConfirmWAL { + private final WriteAheadLog log; + private final Function> commitHandle; + private final List appendListeners = new CopyOnWriteArrayList<>(); + + public ConfirmWAL(WriteAheadLog log, Function> commitHandle) { + this.log = log; + this.commitHandle = commitHandle; + } + + public RecordOffset confirmOffset() { + return log.confirmOffset(); + } + + public String uri() { + return log.uri(); + } + + /** + * Commit with lazy timeout. + * If in [0, lazyLingerMs), there is no other commit happened, then trigger a new commit. + * @param lazyLingerMs lazy linger milliseconds. + */ + public CompletableFuture commit(long lazyLingerMs, boolean awaitTrim) { + return commitHandle.apply(new LazyCommit(lazyLingerMs, awaitTrim)); + } + + public CompletableFuture commit(long lazyLingerMs) { + return commit(lazyLingerMs, true); + } + + public ListenerHandle addAppendListener(AppendListener listener) { + appendListeners.add(listener); + return () -> appendListeners.remove(listener); + } + + public void onAppend(StreamRecordBatch record, RecordOffset recordOffset, RecordOffset nextOffset) { + for (AppendListener listener : appendListeners) { + listener.onAppend(record, recordOffset, nextOffset); + } + } + + public interface AppendListener { + void onAppend(StreamRecordBatch record, RecordOffset recordOffset, RecordOffset nextOffset); + } + + public interface ListenerHandle { + void close(); + } + +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/Constants.java b/s3stream/src/main/java/com/automq/stream/s3/Constants.java index 6ec45fa06a..3961079bf8 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/Constants.java +++ b/s3stream/src/main/java/com/automq/stream/s3/Constants.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/DataBlockIndex.java b/s3stream/src/main/java/com/automq/stream/s3/DataBlockIndex.java index 1815da04f7..f8da9e153d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/DataBlockIndex.java +++ b/s3stream/src/main/java/com/automq/stream/s3/DataBlockIndex.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/DeltaWALUploadTask.java b/s3stream/src/main/java/com/automq/stream/s3/DefaultUploadWriteAheadLogTask.java similarity index 85% rename from s3stream/src/main/java/com/automq/stream/s3/DeltaWALUploadTask.java rename to s3stream/src/main/java/com/automq/stream/s3/DefaultUploadWriteAheadLogTask.java index ffb538eecd..28bc308705 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/DeltaWALUploadTask.java +++ b/s3stream/src/main/java/com/automq/stream/s3/DefaultUploadWriteAheadLogTask.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -36,8 +44,8 @@ import static com.automq.stream.s3.metadata.ObjectUtils.NOOP_OBJECT_ID; -public class DeltaWALUploadTask { - private static final Logger LOGGER = LoggerFactory.getLogger(DeltaWALUploadTask.class); +public class DefaultUploadWriteAheadLogTask implements UploadWriteAheadLogTask { + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultUploadWriteAheadLogTask.class); final boolean forceSplit; private final Logger s3ObjectLogger; private final Map> streamRecordsMap; @@ -57,7 +65,7 @@ public class DeltaWALUploadTask { private volatile CommitStreamSetObjectRequest commitStreamSetObjectRequest; private volatile boolean burst = false; - public DeltaWALUploadTask(Config config, Map> streamRecordsMap, + public DefaultUploadWriteAheadLogTask(Config config, Map> streamRecordsMap, ObjectManager objectManager, ObjectStorage objectStorage, ExecutorService executor, boolean forceSplit, double rate) { this.s3ObjectLogger = S3ObjectLogger.logger(String.format("[DeltaWALUploadTask id=%d] ", config.nodeId())); @@ -77,6 +85,7 @@ public static Builder builder() { return new Builder(); } + @Override public CompletableFuture prepare() { startTimestamp = System.currentTimeMillis(); if (forceSplit) { @@ -96,6 +105,7 @@ public CompletableFuture prepare() { /** * bypass the uploadTask rateLimit to make the task finish as fast as possible. */ + @Override public void burst() { if (this.burst) { return; @@ -113,12 +123,13 @@ private CompletableFuture acquireLimiter(int size) { return limiter.acquire(size); } + @Override public CompletableFuture upload() { prepareCf.thenAcceptAsync(objectId -> FutureUtil.exec(() -> upload0(objectId), uploadCf, LOGGER, "upload"), executor); return uploadCf; } - private void upload0(long objectId) { + void upload0(long objectId) { uploadTimestamp = System.currentTimeMillis(); List streamIds = new ArrayList<>(streamRecordsMap.keySet()); Collections.sort(streamIds); @@ -169,21 +180,27 @@ private void upload0(long objectId) { }); } + @Override public CompletableFuture commit() { return uploadCf.thenCompose(request -> { commitTimestamp = System.currentTimeMillis(); return objectManager.commitStreamSetObject(request).thenAccept(resp -> { long now = System.currentTimeMillis(); - LOGGER.info("Upload delta WAL finished, cost {}ms, prepare {}ms, upload {}ms, commit {}ms, rate limiter {}bytes/s; object id: {}, object size: {}bytes, stream ranges count: {}, stream objects count: {}", + long streamSetObjectSize = request.getObjectSize(); + long streamObjectsSize = request.getStreamObjects().stream().map(StreamObject::getObjectSize).reduce(0L, Long::sum); + long totalSize = streamSetObjectSize + streamObjectsSize; + LOGGER.info("Upload delta WAL finished, cost {}ms, prepare {}ms, upload {}ms, commit {}ms, rate limiter {}bytes/s; object id: {}, object size: {}bytes, stream ranges count: {}, size: {}bytes, stream objects count: {}, size: {}bytes", now - startTimestamp, uploadTimestamp - startTimestamp, commitTimestamp - uploadTimestamp, now - commitTimestamp, (int) rate, request.getObjectId(), - request.getObjectSize(), + totalSize, request.getStreamRanges().size(), - request.getStreamObjects().size() + streamSetObjectSize, + request.getStreamObjects().size(), + streamObjectsSize ); s3ObjectLogger.info("[UPLOAD_WAL] {}", request); }).whenComplete((nil, ex) -> limiter.close()); @@ -256,7 +273,7 @@ public Builder rate(double rate) { return this; } - public DeltaWALUploadTask build() { + public DefaultUploadWriteAheadLogTask build() { if (forceSplit == null) { boolean forceSplit = streamRecordsMap.size() == 1; if (!forceSplit) { @@ -271,7 +288,7 @@ public DeltaWALUploadTask build() { } this.forceSplit = forceSplit; } - return new DeltaWALUploadTask(config, streamRecordsMap, objectManager, objectStorage, executor, forceSplit, rate); + return new DefaultUploadWriteAheadLogTask(config, streamRecordsMap, objectManager, objectStorage, executor, forceSplit, rate); } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java b/s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java index d1ccdc83fe..f840690c60 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/ObjectReader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -18,7 +26,6 @@ import com.automq.stream.s3.network.ThrottleStrategy; import com.automq.stream.s3.objects.ObjectAttributes; import com.automq.stream.s3.operator.ObjectStorage; -import com.automq.stream.s3.operator.ObjectStorage.ReadOptions; import com.automq.stream.utils.CloseableIterator; import com.automq.stream.utils.biniarysearch.IndexBlockOrderedBytes; @@ -33,6 +40,7 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import io.netty.buffer.ByteBuf; @@ -100,6 +108,7 @@ class DefaultObjectReader implements ObjectReader { private CompletableFuture basicObjectInfoCf; private CompletableFuture sizeCf; private final AtomicInteger refCount = new AtomicInteger(1); + private final AtomicBoolean isShutdown = new AtomicBoolean(false); public DefaultObjectReader(S3ObjectMetadata metadata, ObjectStorage objectStorage) { this.metadata = metadata; @@ -116,6 +125,9 @@ public String objectKey() { } public synchronized CompletableFuture basicObjectInfo() { + if (isShutdown.get()) { + return CompletableFuture.failedFuture(new IllegalStateException("ObjectReader is already shutdown")); + } if (basicObjectInfoCf == null) { this.basicObjectInfoCf = new CompletableFuture<>(); asyncGetBasicObjectInfo(); @@ -147,7 +159,7 @@ public CompletableFuture read(ReadOptions readOptions, DataBlock } void asyncGetBasicObjectInfo() { - int guessIndexBlockSize = 1024 + (int) (metadata.objectSize() / (1024 * 1024 /* 1MB */) * 36 /* index unit size*/); + int guessIndexBlockSize = 8192 + (int) (metadata.objectSize() / (1024 * 1024 /* 1MB */) * 36 /* index unit size*/); asyncGetBasicObjectInfo0(Math.max(0, metadata.objectSize() - guessIndexBlockSize), true); } @@ -199,6 +211,9 @@ public synchronized CompletableFuture size() { } public synchronized void close0() { + if (!isShutdown.compareAndSet(false, true)) { + return; + } if (basicObjectInfoCf != null) { basicObjectInfoCf.thenAccept(BasicObjectInfo::close); } @@ -206,17 +221,14 @@ public synchronized void close0() { @Override public boolean equals(Object o) { - if (this == o) - return true; - if (o == null || getClass() != o.getClass()) - return false; - DefaultObjectReader reader = (DefaultObjectReader) o; - return Objects.equals(metadata.objectId(), reader.metadata.objectId()); + // NOTE: DO NOT OVERRIDE THIS + return super.equals(o); } @Override public int hashCode() { - return Objects.hash(metadata.objectId()); + // NOTE: DO NOT OVERRIDE THIS + return super.hashCode(); } } @@ -566,6 +578,10 @@ private static int check(ByteBuf buf) { } public CloseableIterator iterator() { + return iterator(true); + } + + public CloseableIterator iterator(boolean copy) { ByteBuf buf = this.buf.duplicate(); AtomicInteger currentBlockRecordCount = new AtomicInteger(0); AtomicInteger remainingRecordCount = new AtomicInteger(recordCount); @@ -587,7 +603,7 @@ public StreamRecordBatch next() { buf.skipBytes(4); } currentBlockRecordCount.decrementAndGet(); - return StreamRecordBatchCodec.duplicateDecode(buf); + return StreamRecordBatch.parse(buf, copy); } @Override diff --git a/s3stream/src/main/java/com/automq/stream/s3/ObjectWriter.java b/s3stream/src/main/java/com/automq/stream/s3/ObjectWriter.java index 9fb10f7f28..3ab54db80d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/ObjectWriter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/ObjectWriter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -40,9 +48,14 @@ public interface ObjectWriter { // TODO: first n bit is the compressed flag byte DATA_BLOCK_DEFAULT_FLAG = 0x02; + static ObjectWriter writer(long objectId, ObjectStorage objectStorage, int blockSizeThreshold, + int partSizeThreshold, ObjectStorage.WriteOptions writeOptions) { + return new DefaultObjectWriter(objectId, objectStorage, blockSizeThreshold, partSizeThreshold, writeOptions); + } + static ObjectWriter writer(long objectId, ObjectStorage objectStorage, int blockSizeThreshold, int partSizeThreshold) { - return new DefaultObjectWriter(objectId, objectStorage, blockSizeThreshold, partSizeThreshold); + return writer(objectId, objectStorage, blockSizeThreshold, partSizeThreshold, new ObjectStorage.WriteOptions()); } static ObjectWriter noop(long objectId) { @@ -92,15 +105,16 @@ class DefaultObjectWriter implements ObjectWriter { * @param objectStorage S3 operator * @param blockSizeThreshold the max size of a block * @param partSizeThreshold the max size of a part. If it is smaller than {@link Writer#MIN_PART_SIZE}, it will be set to {@link Writer#MIN_PART_SIZE}. + * @param writeOptions the object storage write options */ public DefaultObjectWriter(long objectId, ObjectStorage objectStorage, int blockSizeThreshold, - int partSizeThreshold) { + int partSizeThreshold, ObjectStorage.WriteOptions writeOptions) { String objectKey = ObjectUtils.genKey(0, objectId); this.blockSizeThreshold = blockSizeThreshold; this.partSizeThreshold = Math.max(Writer.MIN_PART_SIZE, partSizeThreshold); waitingUploadBlocks = new LinkedList<>(); completedBlocks = new LinkedList<>(); - writer = objectStorage.writer(ObjectStorage.WriteOptions.DEFAULT, objectKey); + writer = objectStorage.writer(writeOptions, objectKey); } public synchronized void write(long streamId, List records) { diff --git a/s3stream/src/main/java/com/automq/stream/s3/S3ObjectLogger.java b/s3stream/src/main/java/com/automq/stream/s3/S3ObjectLogger.java index a622505c5d..ad6a880772 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/S3ObjectLogger.java +++ b/s3stream/src/main/java/com/automq/stream/s3/S3ObjectLogger.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/S3Storage.java b/s3stream/src/main/java/com/automq/stream/s3/S3Storage.java index f49161726e..cd234220c8 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/S3Storage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/S3Storage.java @@ -1,21 +1,32 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; +import com.automq.stream.Context; +import com.automq.stream.api.LinkRecordDecoder; import com.automq.stream.api.exceptions.FastReadFailFastException; import com.automq.stream.s3.cache.CacheAccessType; import com.automq.stream.s3.cache.LogCache; import com.automq.stream.s3.cache.ReadDataBlock; import com.automq.stream.s3.cache.S3BlockCache; +import com.automq.stream.s3.cache.SnapshotReadCache; import com.automq.stream.s3.context.AppendContext; import com.automq.stream.s3.context.FetchContext; import com.automq.stream.s3.failover.Failover; @@ -30,23 +41,28 @@ import com.automq.stream.s3.streams.StreamManager; import com.automq.stream.s3.trace.context.TraceContext; import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.RecordOffset; import com.automq.stream.s3.wal.RecoverResult; import com.automq.stream.s3.wal.WriteAheadLog; import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.RuntimeIOException; +import com.automq.stream.utils.ExceptionUtil; import com.automq.stream.utils.FutureTicker; import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Systems; import com.automq.stream.utils.ThreadUtils; import com.automq.stream.utils.Threads; +import com.automq.stream.utils.threads.EventLoop; +import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; import java.util.ArrayList; -import java.util.Collections; +import java.util.Collection; import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.LinkedList; import java.util.List; @@ -55,20 +71,20 @@ import java.util.Optional; import java.util.PriorityQueue; import java.util.Queue; +import java.util.Set; +import java.util.concurrent.BlockingQueue; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantLock; -import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -83,25 +99,44 @@ public class S3Storage implements Storage { private static final FastReadFailFastException FAST_READ_FAIL_FAST_EXCEPTION = new FastReadFailFastException(); private static final int NUM_STREAM_CALLBACK_LOCKS = 128; - private final long maxDeltaWALCacheSize; - private final Config config; + + private static LinkRecordDecoder linkRecordDecoder = LinkRecordDecoder.NOOP; + + public static void setLinkRecordDecoder( + LinkRecordDecoder linkRecordDecoder) { + S3Storage.linkRecordDecoder = linkRecordDecoder; + } + + public static LinkRecordDecoder getLinkRecordDecoder() { + return linkRecordDecoder; + } + + protected final Config config; private final WriteAheadLog deltaWAL; + private final ConfirmWAL confirmWAL; /** * WAL log cache */ private final LogCache deltaWALCache; - /** - * WAL out of order callback sequencer. {@link #streamCallbackLocks} will ensure the memory safety. - */ - private final WALCallbackSequencer callbackSequencer = new WALCallbackSequencer(); - private final WALConfirmOffsetCalculator confirmOffsetCalculator = new WALConfirmOffsetCalculator(); + private final LogCache snapshotReadCache; private final Queue walPrepareQueue = new LinkedList<>(); private final Queue walCommitQueue = new LinkedList<>(); private final List inflightWALUploadTasks = new CopyOnWriteArrayList<>(); + private BlockingQueue lazyUploadQueue = new LinkedBlockingQueue<>(); + + /** + * A lock to ensure only one thread can trigger {@link #forceUpload()} in {@link #maybeForceUpload()} + */ + private final AtomicBoolean forceUploadScheduled = new AtomicBoolean(); + /** + * A lock to ensure only one thread can trigger {@link #forceUpload()} in {@link #forceUploadCallback()} + */ + private final AtomicBoolean needForceUpload = new AtomicBoolean(); private final ScheduledExecutorService backgroundExecutor = Threads.newSingleThreadScheduledExecutor( ThreadUtils.createThreadFactory("s3-storage-background", true), LOGGER); private final ExecutorService uploadWALExecutor = Threads.newFixedThreadPoolWithMonitor( 4, "s3-storage-upload-wal", true, LOGGER); + private final DelayTrim delayTrim; /** * A ticker used for batching force upload WAL. * @@ -110,17 +145,20 @@ public class S3Storage implements Storage { private final FutureTicker forceUploadTicker = new FutureTicker(100, TimeUnit.MILLISECONDS, backgroundExecutor); private final Queue backoffRecords = new LinkedBlockingQueue<>(); private final ScheduledFuture drainBackoffTask; - private final StreamManager streamManager; - private final ObjectManager objectManager; - private final ObjectStorage objectStorage; - private final S3BlockCache blockCache; - private final StorageFailureHandler storageFailureHandler; + protected final StreamManager streamManager; + protected final ObjectManager objectManager; + protected final ObjectStorage objectStorage; + protected final S3BlockCache blockCache; + protected final StorageFailureHandler storageFailureHandler; /** * Stream callback locks. Used to ensure the stream callbacks will not be called concurrently. * * @see #handleAppendCallback */ private final Lock[] streamCallbackLocks = IntStream.range(0, NUM_STREAM_CALLBACK_LOCKS).mapToObj(i -> new ReentrantLock()).toArray(Lock[]::new); + private final EventLoop[] callbackExecutors = IntStream.range(0, Systems.CPU_CORES).mapToObj(i -> new EventLoop("AUTOMQ_S3STREAM_APPEND_CALLBACK-" + i)) + .toArray(EventLoop[]::new); + private long lastLogTimestamp = 0L; private volatile double maxDataWriteRate = 0.0; @@ -130,10 +168,28 @@ public class S3Storage implements Storage { public S3Storage(Config config, WriteAheadLog deltaWAL, StreamManager streamManager, ObjectManager objectManager, S3BlockCache blockCache, ObjectStorage objectStorage, StorageFailureHandler storageFailureHandler) { this.config = config; - this.maxDeltaWALCacheSize = config.walCacheSize(); this.deltaWAL = deltaWAL; this.blockCache = blockCache; - this.deltaWALCache = new LogCache(config.walCacheSize(), config.walUploadThreshold(), config.maxStreamNumPerStreamSetObject()); + long deltaWALCacheSize = config.walCacheSize(); + long snapshotReadCacheSize = 0; + if (config.snapshotReadEnable()) { + deltaWALCacheSize = Math.max(config.walCacheSize() / 3, 10L * 1024 * 1024); + snapshotReadCacheSize = Math.max(config.walCacheSize() / 3 * 2, 10L * 1024 * 1024); + delayTrim = new DelayTrim(TimeUnit.SECONDS.toMillis(30)); + } else { + delayTrim = new DelayTrim(0); + } + // Adjust the walUploadThreshold to be less than 2/5 of deltaWALCacheSize to avoid the upload speed being slower than the append speed. + long walUploadThreadhold = Math.min(deltaWALCacheSize * 2 / 5, config.walUploadThreshold()); + if (walUploadThreadhold != config.walUploadThreshold()) { + LOGGER.info("The configured walUploadThreshold {} is too large, adjust to {}", config.walUploadThreshold(), walUploadThreadhold); + } + this.deltaWALCache = new LogCache(deltaWALCacheSize, walUploadThreadhold, config.maxStreamNumPerStreamSetObject()); + this.snapshotReadCache = new LogCache(snapshotReadCacheSize, Math.max(snapshotReadCacheSize / 6, 1)); + S3StreamMetricsManager.registerDeltaWalCacheSizeSupplier(() -> deltaWALCache.size() + snapshotReadCache.size()); + Context.instance().snapshotReadCache(new SnapshotReadCache(streamManager, snapshotReadCache, objectStorage, linkRecordDecoder)); + this.confirmWAL = new ConfirmWAL(deltaWAL, lazyCommit -> lazyUpload(lazyCommit)); + Context.instance().confirmWAL(this.confirmWAL); this.streamManager = streamManager; this.objectManager = objectManager; this.objectStorage = objectStorage; @@ -141,22 +197,14 @@ public S3Storage(Config config, WriteAheadLog deltaWAL, StreamManager streamMana this.drainBackoffTask = this.backgroundExecutor.scheduleWithFixedDelay(this::tryDrainBackoffRecords, 100, 100, TimeUnit.MILLISECONDS); S3StreamMetricsManager.registerInflightWALUploadTasksCountSupplier(this.inflightWALUploadTasks::size); S3StreamMetricsManager.registerDeltaWalPendingUploadBytesSupplier(this.pendingUploadBytes::get); - } - - /** - * Only for test. - */ - static LogCache.LogCacheBlock recoverContinuousRecords(Iterator it, - List openingStreams) { - InnerRecoverResult result = recoverContinuousRecords(it, openingStreams, LOGGER); - result.firstException().ifPresent(e -> { - throw e; - }); - return result.cacheBlock; + if (config.walUploadIntervalMs() > 0) { + this.backgroundExecutor.scheduleWithFixedDelay(this::maybeForceUpload, config.walUploadIntervalMs(), config.walUploadIntervalMs(), TimeUnit.MILLISECONDS); + } } /** * Recover continuous records in each stream from the WAL, and put them into the returned {@link LogCache.LogCacheBlock}. + *

    * It will filter out *

      *
    • the records that are not in the opening streams
      • @@ -180,28 +228,143 @@ static LogCache.LogCacheBlock recoverContinuousRecords(Iterator i *
    • The record 10 and 11 are discarded because they are not continuous (10 is not 7, the next offset of 6)
      • *
    • The record 5 and 4 are reordered because they are out of order, and we handle this bug here
      • *
    + *

    + * It will return when any of the following conditions is met: + *

      + *
    • all the records in the WAL have been recovered
      • + *
    • the cache block is full
      • + *
    + * Visible for testing. + * + * @param it WAL recover iterator + * @param openingStreamEndOffsets the end offset of each opening stream + * @param maxCacheSize the max size of the returned {@link RecoveryBlockResult#cacheBlock} + * @param logger logger */ - static InnerRecoverResult recoverContinuousRecords(Iterator it, - List openingStreams, Logger logger) { - Map openingStreamEndOffsets = openingStreams.stream().collect(Collectors.toMap(StreamMetadata::streamId, StreamMetadata::endOffset)); - LogCache.LogCacheBlock cacheBlock = new LogCache.LogCacheBlock(1024L * 1024 * 1024); + static RecoveryBlockResult recoverContinuousRecords( + Iterator it, + Map openingStreamEndOffsets, + long maxCacheSize, + Logger logger + ) { + RecordOffset logEndOffset = null; Map streamNextOffsets = new HashMap<>(); Map> streamDiscontinuousRecords = new HashMap<>(); - long logEndOffset = recoverContinuousRecords(it, openingStreamEndOffsets, streamNextOffsets, streamDiscontinuousRecords, cacheBlock, logger); - // release all discontinuous records. - streamDiscontinuousRecords.values().forEach(queue -> { - if (queue.isEmpty()) { - return; + LogCache.LogCacheBlock cacheBlock = new LogCache.LogCacheBlock(maxCacheSize); + + boolean first = true; + try { + while (it.hasNext() && !cacheBlock.isFull()) { + RecoverResult recoverResult = it.next(); + logEndOffset = recoverResult.recordOffset(); + if (first) { + LOGGER.info("recover start offset {}", logEndOffset); + first = false; + } + StreamRecordBatch streamRecordBatch = recoverResult.record(); + processRecoveredRecord(streamRecordBatch, openingStreamEndOffsets, streamDiscontinuousRecords, cacheBlock, streamNextOffsets, logger); } - logger.info("drop discontinuous records, records={}", queue); - queue.forEach(StreamRecordBatch::release); - }); + } catch (Throwable e) { + // {@link RuntimeIOException} may be thrown by {@code it.next()} + releaseAllRecords(streamDiscontinuousRecords.values()); + releaseAllRecords(cacheBlock.records().values()); + throw e; + } + if (logEndOffset != null) { + cacheBlock.lastRecordOffset(logEndOffset); + } + + releaseDiscontinuousRecords(streamDiscontinuousRecords, logger); + RecoveryBlockResult rst = filterOutInvalidStreams(cacheBlock, openingStreamEndOffsets); + return decodeLinkRecord(rst); + } + + /** + * Processes recovered stream records. Caches continuous ones or queues discontinuous based on offset order. + * + * @param streamRecordBatch the recovered record batch to process + * @param openingStreamEndOffsets the end offsets of each opening stream + * @param streamDiscontinuousRecords the out-of-order records of each stream (to be filled) + * @param cacheBlock the cache block (to be filled) + * @param streamNextOffsets the next offsets of each stream (to be updated) + * @param logger logger + */ + private static void processRecoveredRecord( + StreamRecordBatch streamRecordBatch, + Map openingStreamEndOffsets, + Map> streamDiscontinuousRecords, + LogCache.LogCacheBlock cacheBlock, + Map streamNextOffsets, + Logger logger + ) { + long streamId = streamRecordBatch.getStreamId(); + + Long openingStreamEndOffset = openingStreamEndOffsets.get(streamId); + if (openingStreamEndOffset == null || openingStreamEndOffset > streamRecordBatch.getBaseOffset()) { + // stream is already safe closed, or the record have been committed, skip it + streamRecordBatch.release(); + return; + } + + Long expectedNextOffset = streamNextOffsets.get(streamId); + Queue discontinuousRecords = streamDiscontinuousRecords.get(streamId); + boolean isContinuous = expectedNextOffset == null || expectedNextOffset == streamRecordBatch.getBaseOffset(); + if (!isContinuous) { + // unexpected record, put it into discontinuous records queue. + if (discontinuousRecords == null) { + discontinuousRecords = new PriorityQueue<>(Comparator.comparingLong(StreamRecordBatch::getBaseOffset)); + streamDiscontinuousRecords.put(streamId, discontinuousRecords); + } + discontinuousRecords.add(streamRecordBatch); + return; + } + // continuous record, put it into cache, and check if there is any historical discontinuous records can be polled. + cacheBlock.put(streamRecordBatch); + expectedNextOffset = maybePollDiscontinuousRecords(streamRecordBatch, cacheBlock, discontinuousRecords, logger); + streamNextOffsets.put(streamId, expectedNextOffset); + } - if (logEndOffset >= 0L) { - cacheBlock.confirmOffset(logEndOffset); + private static long maybePollDiscontinuousRecords( + StreamRecordBatch streamRecordBatch, + LogCache.LogCacheBlock cacheBlock, + Queue discontinuousRecords, + Logger logger + ) { + long expectedNextOffset = streamRecordBatch.getLastOffset(); + if (discontinuousRecords == null) { + return expectedNextOffset; } + // check and poll historical discontinuous records. + while (!discontinuousRecords.isEmpty()) { + StreamRecordBatch peek = discontinuousRecords.peek(); + if (peek.getBaseOffset() != expectedNextOffset) { + break; + } + // should never happen, log it. + logger.error("[BUG] recover an out of order record, streamId={}, expectedNextOffset={}, record={}", streamRecordBatch.getStreamId(), expectedNextOffset, peek); + discontinuousRecords.poll(); + cacheBlock.put(peek); + expectedNextOffset = peek.getLastOffset(); + } + return expectedNextOffset; + } + + private static void releaseDiscontinuousRecords(Map> streamDiscontinuousRecords, + Logger logger) { + streamDiscontinuousRecords.values().stream() + .filter(q -> !q.isEmpty()) + .peek(q -> logger.info("drop discontinuous records, records={}", q)) + .forEach(S3Storage::releaseRecords); + } + + /** + * Filter out invalid streams (the recovered start offset mismatches the stream end offset from controller) from the cache block if there are any. + */ + private static RecoveryBlockResult filterOutInvalidStreams(LogCache.LogCacheBlock cacheBlock, + Map openingStreamEndOffsets) { + Set invalidStreams = new HashSet<>(); + List exceptions = new ArrayList<>(); - InnerRecoverResult result = new InnerRecoverResult(); cacheBlock.records().forEach((streamId, records) -> { if (!records.isEmpty()) { long startOffset = records.get(0).getBaseOffset(); @@ -210,110 +373,65 @@ static InnerRecoverResult recoverContinuousRecords(Iterator it, RuntimeException exception = new IllegalStateException(String.format("[BUG] WAL data may lost, streamId %d endOffset=%d from controller, " + "but WAL recovered records startOffset=%s", streamId, expectedStartOffset, startOffset)); LOGGER.error("invalid stream records", exception); - result.invalidStreams.put(streamId, exception); + invalidStreams.add(streamId); + exceptions.add(exception); } } }); - if (result.invalidStreams.isEmpty()) { - result.cacheBlock = cacheBlock; - } else { - // re-new a cache block and put all valid records into it. - LogCache.LogCacheBlock newCacheBlock = new LogCache.LogCacheBlock(1024L * 1024 * 1024); - cacheBlock.records().forEach((streamId, records) -> { - if (!result.invalidStreams.containsKey(streamId)) { - records.forEach(newCacheBlock::put); - } else { - // release invalid records. - records.forEach(StreamRecordBatch::release); - } - }); - result.cacheBlock = newCacheBlock; + + if (invalidStreams.isEmpty()) { + return new RecoveryBlockResult(cacheBlock, null); } - return result; + // Only streams not in invalidStreams should be uploaded and closed, + // so re-new a cache block and put only valid records into it, and release all invalid records. + LogCache.LogCacheBlock newCacheBlock = new LogCache.LogCacheBlock(1024L * 1024 * 1024); + cacheBlock.records().forEach((streamId, records) -> { + if (!invalidStreams.contains(streamId)) { + records.forEach(newCacheBlock::put); + } else { + // release invalid records. + releaseRecords(records); + } + }); + return new RecoveryBlockResult(newCacheBlock, ExceptionUtil.combine(exceptions)); } - private static long recoverContinuousRecords(Iterator it, Map openingStreamEndOffsets, - Map streamNextOffsets, Map> streamDiscontinuousRecords, - LogCache.LogCacheBlock cacheBlock, Logger logger) { - try { - return recoverContinuousRecords0(it, openingStreamEndOffsets, streamNextOffsets, streamDiscontinuousRecords, cacheBlock, logger); - } catch (Throwable e) { - streamDiscontinuousRecords.values().forEach(queue -> queue.forEach(StreamRecordBatch::release)); - cacheBlock.records().forEach((streamId, records) -> records.forEach(StreamRecordBatch::release)); - throw e; - } + private static void releaseAllRecords(Collection> allRecords) { + allRecords.forEach(S3Storage::releaseRecords); } - /** - * Recover continuous records in each stream from the WAL, and put them into the returned {@link LogCache.LogCacheBlock}. - * - * @param it WAL recover iterator - * @param openingStreamEndOffsets the end offset of each opening stream - * @param streamNextOffsets the next offset of each stream (to be filled) - * @param streamDiscontinuousRecords the out-of-order records of each stream (to be filled) - * @param cacheBlock the cache block (to be filled) - * @return the end offset of the last record recovered - * @throws RuntimeIOException if any IO error occurs during recover from Block WAL - */ - private static long recoverContinuousRecords0(Iterator it, - Map openingStreamEndOffsets, - Map streamNextOffsets, - Map> streamDiscontinuousRecords, - LogCache.LogCacheBlock cacheBlock, - Logger logger) throws RuntimeIOException { - long logEndOffset = -1L; - while (it.hasNext()) { - RecoverResult recoverResult = it.next(); - logEndOffset = recoverResult.recordOffset(); - ByteBuf recordBuf = recoverResult.record().duplicate(); - StreamRecordBatch streamRecordBatch = StreamRecordBatchCodec.decode(recordBuf); - long streamId = streamRecordBatch.getStreamId(); - Long openingStreamEndOffset = openingStreamEndOffsets.get(streamId); - if (openingStreamEndOffset == null) { - // stream is already safe closed. so skip the stream records. - recordBuf.release(); - continue; - } - if (streamRecordBatch.getBaseOffset() < openingStreamEndOffset) { - // filter committed records. - recordBuf.release(); - continue; - } + private static void releaseRecords(Collection records) { + records.forEach(StreamRecordBatch::release); + } - Long expectNextOffset = streamNextOffsets.get(streamId); - Queue discontinuousRecords = streamDiscontinuousRecords.get(streamId); - if (expectNextOffset == null || expectNextOffset == streamRecordBatch.getBaseOffset()) { - // continuous record, put it into cache. - cacheBlock.put(streamRecordBatch); - expectNextOffset = streamRecordBatch.getLastOffset(); - // check if there are some out of order records in the queue. - if (discontinuousRecords != null) { - while (!discontinuousRecords.isEmpty()) { - StreamRecordBatch peek = discontinuousRecords.peek(); - if (peek.getBaseOffset() == expectNextOffset) { - // should never happen, log it. - logger.error("[BUG] recover an out of order record, streamId={}, expectNextOffset={}, record={}", streamId, expectNextOffset, peek); - cacheBlock.put(peek); - discontinuousRecords.poll(); - expectNextOffset = peek.getLastOffset(); - } else { - break; - } - } - } - // update next offset. - streamNextOffsets.put(streamRecordBatch.getStreamId(), expectNextOffset); - } else { - // unexpected record, put it into discontinuous records queue. - if (discontinuousRecords == null) { - discontinuousRecords = new PriorityQueue<>(Comparator.comparingLong(StreamRecordBatch::getBaseOffset)); - streamDiscontinuousRecords.put(streamId, discontinuousRecords); + private static RecoveryBlockResult decodeLinkRecord(RecoveryBlockResult recoverBlockRst) { + LogCache.LogCacheBlock cacheBlock = recoverBlockRst.cacheBlock; + int size = 0; + for (List l : cacheBlock.records().values()) { + size += l.size(); + } + List> futures = new ArrayList<>(size); + for (Map.Entry> entry : cacheBlock.records().entrySet()) { + List records = entry.getValue(); + for (int i = 0; i < records.size(); i++) { + StreamRecordBatch record = records.get(i); + if (record.getCount() >= 0) { + continue; } - discontinuousRecords.add(streamRecordBatch); + int finalI = i; + futures.add(linkRecordDecoder.decode(record).thenAccept(r -> { + records.set(finalI, r); + })); } } - return logEndOffset; + try { + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(); + return recoverBlockRst; + } catch (Throwable ex) { + releaseAllRecords(cacheBlock.records().values()); + throw new RuntimeException(ex); + } } @Override @@ -350,42 +468,63 @@ public void recover(WriteAheadLog deltaWAL, StreamManager streamManager, ObjectM * Note: {@link WriteAheadLog#start()} should be called before this method. */ void recover0(WriteAheadLog deltaWAL, StreamManager streamManager, ObjectManager objectManager, - Logger logger) throws Throwable { + Logger logger) throws InterruptedException, ExecutionException { List streams = streamManager.getOpeningStreams().get(); + Map streamEndOffsets = streams.stream().collect(Collectors.toMap(StreamMetadata::streamId, StreamMetadata::endOffset)); + Iterator iterator = deltaWAL.recover(); + + LogCache.LogCacheBlock cacheBlock; + List exceptions = new ArrayList<>(); + do { + RecoveryBlockResult result = recoverContinuousRecords(iterator, streamEndOffsets, 1 << 29, logger); + cacheBlock = result.cacheBlock; + Optional.ofNullable(result.exception).ifPresent(exceptions::add); + updateStreamEndOffsets(cacheBlock, streamEndOffsets); + uploadRecoveredRecords(objectManager, cacheBlock, logger); + } + while (cacheBlock.isFull()); - InnerRecoverResult recoverResult = recoverContinuousRecords(deltaWAL.recover(), streams, logger); - LogCache.LogCacheBlock cacheBlock = recoverResult.cacheBlock; + deltaWAL.reset().get(); + closeStreams(streamManager, streams, streamEndOffsets, logger); + + // fail it if there is any invalid stream. + if (!exceptions.isEmpty()) { + throw ExceptionUtil.combine(exceptions); + } + } - Map streamEndOffsets = new HashMap<>(); + private static void updateStreamEndOffsets(LogCache.LogCacheBlock cacheBlock, Map streamEndOffsets) { cacheBlock.records().forEach((streamId, records) -> { if (!records.isEmpty()) { streamEndOffsets.put(streamId, records.get(records.size() - 1).getLastOffset()); } }); + } + private void uploadRecoveredRecords(ObjectManager objectManager, LogCache.LogCacheBlock cacheBlock, Logger logger) + throws InterruptedException, ExecutionException { if (cacheBlock.size() != 0) { logger.info("try recover from crash, recover records bytes size {}", cacheBlock.size()); - DeltaWALUploadTask task = DeltaWALUploadTask.builder().config(config).streamRecordsMap(cacheBlock.records()) - .objectManager(objectManager).objectStorage(objectStorage).executor(uploadWALExecutor).build(); - task.prepare().thenCompose(nil -> task.upload()).thenCompose(nil -> task.commit()).get(); - cacheBlock.records().forEach((streamId, records) -> records.forEach(StreamRecordBatch::release)); + try { + UploadWriteAheadLogTask task = newUploadWriteAheadLogTask(cacheBlock.records(), objectManager, Long.MAX_VALUE); + task.prepare().thenCompose(nil -> task.upload()).thenCompose(nil -> task.commit()).get(); + } finally { + releaseAllRecords(cacheBlock.records().values()); + } } - deltaWAL.reset().get(); + } + + private static void closeStreams(StreamManager streamManager, List streams, + Map streamEndOffsets, Logger logger) throws InterruptedException, ExecutionException { for (StreamMetadata stream : streams) { long newEndOffset = streamEndOffsets.getOrDefault(stream.streamId(), stream.endOffset()); logger.info("recover try close stream {} with new end offset {}", stream, newEndOffset); } CompletableFuture.allOf( - streams - .stream() + streams.stream() .map(s -> streamManager.closeStream(s.streamId(), s.epoch())) .toArray(CompletableFuture[]::new) ).get(); - - // fail it if there is any invalid stream. - recoverResult.firstException().ifPresent(e -> { - throw e; - }); } @Override @@ -394,15 +533,11 @@ public void shutdown() { for (WalWriteRequest request : backoffRecords) { request.cf.completeExceptionally(new IOException("S3Storage is shutdown")); } + suppress(() -> delayTrim.close(), LOGGER); deltaWAL.shutdownGracefully(); - backgroundExecutor.shutdown(); - try { - if (backgroundExecutor.awaitTermination(10, TimeUnit.SECONDS)) { - LOGGER.warn("await backgroundExecutor timeout 10s"); - } - } catch (InterruptedException e) { - backgroundExecutor.shutdownNow(); - LOGGER.warn("await backgroundExecutor close fail", e); + ThreadUtils.shutdownExecutor(backgroundExecutor, 10, TimeUnit.SECONDS, LOGGER); + for (EventLoop executor : callbackExecutors) { + executor.shutdownGracefully(); } } @@ -411,10 +546,7 @@ public void shutdown() { public CompletableFuture append(AppendContext context, StreamRecordBatch streamRecord) { final long startTime = System.nanoTime(); CompletableFuture cf = new CompletableFuture<>(); - // encoded before append to free heap ByteBuf. - streamRecord.encoded(); - WalWriteRequest writeRequest = new WalWriteRequest(streamRecord, -1L, cf, context); - handleAppendRequest(writeRequest); + WalWriteRequest writeRequest = new WalWriteRequest(streamRecord, null, cf, context); append0(context, writeRequest, false); return cf.whenComplete((nil, ex) -> { streamRecord.release(); @@ -440,27 +572,27 @@ public boolean append0(AppendContext context, WalWriteRequest request, boolean f } StorageOperationStats.getInstance().appendLogCacheFullStats.record(0L); if (System.currentTimeMillis() - lastLogTimestamp > 1000L) { - LOGGER.warn("[BACKOFF] log cache size {} is larger than {}", deltaWALCache.size(), maxDeltaWALCacheSize); + LOGGER.warn("[BACKOFF] log cache size {} is larger than {}", deltaWALCache.size(), deltaWALCache.capacity()); lastLogTimestamp = System.currentTimeMillis(); } return true; } - AppendResult appendResult; + CompletableFuture appendCf; try { try { - StreamRecordBatch streamRecord = request.record; - streamRecord.retain(); - Lock lock = confirmOffsetCalculator.addLock(); - lock.lock(); - try { - appendResult = deltaWAL.append(new TraceContext(context), streamRecord.encoded()); - } finally { - lock.unlock(); + if (context.linkRecord() == null) { + StreamRecordBatch streamRecord = request.record; + streamRecord.retain(); + appendCf = deltaWAL.append(new TraceContext(context), streamRecord); + } else { + StreamRecordBatch record = request.record; + StreamRecordBatch linkStreamRecord = toLinkRecord(record, context.linkRecord().retainedSlice()); + appendCf = deltaWAL.append(new TraceContext(context), linkStreamRecord); } + } catch (OverCapacityException e) { // the WAL write data align with block, 'WAL is full but LogCacheBlock is not full' may happen. - confirmOffsetCalculator.update(); - forceUpload(LogCache.MATCH_ALL_STREAMS); + maybeForceUpload(); if (!fromBackoff) { backoffRecords.offer(request); } @@ -470,30 +602,35 @@ public boolean append0(AppendContext context, WalWriteRequest request, boolean f } return true; } - long recordOffset = appendResult.recordOffset(); - if (recordOffset >= 0) { - request.offset = recordOffset; - confirmOffsetCalculator.add(request); - } } catch (Throwable e) { LOGGER.error("[UNEXPECTED] append WAL fail", e); request.cf.completeExceptionally(e); return false; } - appendResult.future().whenComplete((nil, ex) -> { + appendCf.thenAccept(rst -> { + request.offset = rst.recordOffset(); + // Execute the ConfirmWAL#append before run callback. + if (request.context.linkRecord() == null) { + this.confirmWAL.onAppend(request.record, rst.recordOffset(), rst.nextOffset()); + } else { + StreamRecordBatch linkRecord = toLinkRecord(request.record, request.context.linkRecord()); + this.confirmWAL.onAppend(linkRecord, rst.recordOffset(), rst.nextOffset()); + linkRecord.release(); + } + handleAppendCallback(request); + }).whenComplete((nil, ex) -> { if (ex != null) { - LOGGER.error("append WAL fail, request {}", request, ex); + LOGGER.error("append WAL fail", ex); storageFailureHandler.handle(ex); return; } - handleAppendCallback(request); }); return false; } @SuppressWarnings("BooleanMethodIsAlwaysInverted") private boolean tryAcquirePermit() { - return deltaWALCache.size() < maxDeltaWALCacheSize; + return deltaWALCache.size() < deltaWALCache.capacity(); } private void tryDrainBackoffRecords() { @@ -528,19 +665,25 @@ public CompletableFuture read(FetchContext context, return cf; } + public LogCache snapshotReadCache() { + return snapshotReadCache; + } + + @SuppressWarnings({"checkstyle:npathcomplexity"}) @WithSpan private CompletableFuture read0(FetchContext context, @SpanAttribute long streamId, @SpanAttribute long startOffset, @SpanAttribute long endOffset, @SpanAttribute int maxBytes) { - List logCacheRecords = deltaWALCache.get(context, streamId, startOffset, endOffset, maxBytes); + LogCache firstCache = context.readOptions().snapshotRead() ? snapshotReadCache : deltaWALCache; + List logCacheRecords = firstCache.get(context, streamId, startOffset, endOffset, maxBytes); if (!logCacheRecords.isEmpty() && logCacheRecords.get(0).getBaseOffset() <= startOffset) { return CompletableFuture.completedFuture(new ReadDataBlock(logCacheRecords, CacheAccessType.DELTA_WAL_CACHE_HIT)); } if (context.readOptions().fastRead()) { // fast read fail fast when need read from block cache. - logCacheRecords.forEach(StreamRecordBatch::release); + releaseRecords(logCacheRecords); logCacheRecords.clear(); return CompletableFuture.failedFuture(FAST_READ_FAIL_FAST_EXCEPTION); } @@ -561,38 +704,117 @@ private CompletableFuture read0(FetchContext context, try { continuousCheck(rst); } catch (IllegalArgumentException e) { - blockCacheRst.getRecords().forEach(StreamRecordBatch::release); + releaseRecords(blockCacheRst.getRecords()); throw e; } if (readIndex < logCacheRecords.size()) { // release unnecessary record - logCacheRecords.subList(readIndex + 1, logCacheRecords.size()).forEach(StreamRecordBatch::release); + releaseRecords(logCacheRecords.subList(readIndex + 1, logCacheRecords.size())); } return new ReadDataBlock(rst, blockCacheRst.getCacheAccessType()); }).whenComplete((rst, ex) -> { if (ex != null) { LOGGER.error("read from block cache failed, stream={}, {}-{}, maxBytes: {}", streamId, startOffset, finalEndOffset, maxBytes, ex); - logCacheRecords.forEach(StreamRecordBatch::release); + releaseRecords(logCacheRecords); } }); return FutureUtil.timeoutWithNewReturn(cf, 2, TimeUnit.MINUTES, () -> { LOGGER.error("[POTENTIAL_BUG] read from block cache timeout, stream={}, [{},{}), maxBytes: {}", streamId, startOffset, finalEndOffset, maxBytes); cf.thenAccept(readDataBlock -> { - readDataBlock.getRecords().forEach(r -> r.release()); + releaseRecords(readDataBlock.getRecords()); }); }); } private void continuousCheck(List records) { - long expectStartOffset = -1L; + long expectedStartOffset = -1L; for (StreamRecordBatch record : records) { - if (expectStartOffset == -1L || record.getBaseOffset() == expectStartOffset) { - expectStartOffset = record.getLastOffset(); + if (expectedStartOffset == -1L || record.getBaseOffset() == expectedStartOffset) { + expectedStartOffset = record.getLastOffset(); } else { - throw new IllegalArgumentException(String.format("Continuous check failed, expect offset: %d," + - " actual: %d, records: %s", expectStartOffset, record.getBaseOffset(), records)); + throw new IllegalArgumentException(String.format("Continuous check failed, expected offset: %d," + + " actual: %d, records: %s", expectedStartOffset, record.getBaseOffset(), records)); + } + } + } + + /** + * Limit the number of inflight force upload tasks to 1 to avoid too many S3 objects. + */ + private void maybeForceUpload() { + if (hasInflightForceUploadTask()) { + // There is already an inflight force upload task, trigger another one later after it completes. + needForceUpload.set(true); + return; + } + if (forceUploadScheduled.compareAndSet(false, true)) { + forceUpload(); + } else { + // There is already a force upload task scheduled, do nothing. + needForceUpload.set(true); + } + } + + private boolean hasInflightForceUploadTask() { + return inflightWALUploadTasks.stream().anyMatch(it -> it.force); + } + + /** + * Commit with lazy timeout. If in [0, lazyLingerMs), there is no other commit happened, then trigger a new commit. + */ + private CompletableFuture lazyUpload(LazyCommit lazyCommit) { + lazyUploadQueue.add(lazyCommit); + backgroundExecutor.schedule(() -> { + if (lazyUploadQueue.contains(lazyCommit)) { + // If the queue does not contain the lazyCommit, it means another commit has happened after this lazyCommit. + if (lazyCommit.lazyLingerMs == 0) { + // If the lazyLingerMs is 0, we need to force upload as soon as possible. + forceUpload(); + } else { + uploadDeltaWAL(); + } } + }, lazyCommit.lazyLingerMs, TimeUnit.MILLISECONDS); + return lazyCommit.awaitTrim ? lazyCommit.trimCf : lazyCommit.commitCf; + } + + private void notifyLazyUpload(List tasks) { + CompletableFuture.allOf(inflightWALUploadTasks.stream().map(t -> t.cf).collect(Collectors.toList()).toArray(new CompletableFuture[0])) + .whenComplete((nil, ex) -> { + for (LazyCommit task : tasks) { + if (ex != null) { + task.commitCf.completeExceptionally(ex); + } else { + task.commitCf.complete(null); + } + } + }); + + CompletableFuture.allOf(inflightWALUploadTasks.stream().map(t -> t.trimCf).collect(Collectors.toList()).toArray(new CompletableFuture[0])) + .whenComplete((nil, ex) -> { + for (LazyCommit task : tasks) { + if (ex != null) { + task.trimCf.completeExceptionally(ex); + } else { + task.trimCf.complete(null); + } + } + }); + } + + private CompletableFuture forceUpload() { + CompletableFuture cf = forceUpload(LogCache.MATCH_ALL_STREAMS); + cf.whenComplete((nil, ignored) -> forceUploadCallback()); + return cf; + } + + private void forceUploadCallback() { + // Reset the force upload flag after the task completes. + forceUploadScheduled.set(false); + if (needForceUpload.compareAndSet(true, false)) { + // Force upload needs to be triggered again. + forceUpload(); } } @@ -612,71 +834,70 @@ public CompletableFuture forceUpload(long streamId) { FutureUtil.propagate(CompletableFuture.allOf(this.inflightWALUploadTasks.stream() .filter(it -> it.cache.containsStream(streamId)) .map(it -> it.cf).toArray(CompletableFuture[]::new)), cf); - if (LogCache.MATCH_ALL_STREAMS != streamId) { - callbackSequencer.tryFree(streamId); - } }); cf.whenComplete((nil, ex) -> StorageOperationStats.getInstance().forceUploadWALCompleteStats.record( TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS))); return cf; } - private void handleAppendRequest(WalWriteRequest request) { - callbackSequencer.before(request); - } - private void handleAppendCallback(WalWriteRequest request) { - suppress(() -> handleAppendCallback0(request), LOGGER); - } - - private void handleAppendCallback0(WalWriteRequest request) { final long startTime = System.nanoTime(); - List waitingAckRequests; - Lock lock = getStreamCallbackLock(request.record.getStreamId()); - lock.lock(); - try { - waitingAckRequests = callbackSequencer.after(request); - waitingAckRequests.forEach(r -> r.record.retain()); - for (WalWriteRequest waitingAckRequest : waitingAckRequests) { - boolean full = deltaWALCache.put(waitingAckRequest.record); - waitingAckRequest.confirmed = true; - if (full) { - // cache block is full, trigger WAL upload. - uploadDeltaWAL(); - } - } - } finally { - lock.unlock(); + request.record.retain(); + boolean full; + synchronized (deltaWALCache) { + // Because LogCacheBlock will use request.offset to execute WAL#trim after being uploaded, + // this cache put order should keep consistence with WAL put order. + full = deltaWALCache.put(request.record); + deltaWALCache.setLastRecordOffset(request.offset); } - for (WalWriteRequest waitingAckRequest : waitingAckRequests) { - waitingAckRequest.cf.complete(null); + if (full) { + // cache block is full, trigger WAL upload. + uploadDeltaWAL(); } - StorageOperationStats.getInstance().appendCallbackStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); + // parallel execute append callback in streamId based executor. + EventLoop executor = callbackExecutors[Math.abs((int) (request.record.getStreamId() % callbackExecutors.length))]; + executor.execute(() -> { + request.cf.complete(null); + StorageOperationStats.getInstance().appendCallbackStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); + }); } private Lock getStreamCallbackLock(long streamId) { return streamCallbackLocks[(int) ((streamId & Long.MAX_VALUE) % NUM_STREAM_CALLBACK_LOCKS)]; } + protected UploadWriteAheadLogTask newUploadWriteAheadLogTask(Map> streamRecordsMap, + ObjectManager objectManager, double rate) { + return DefaultUploadWriteAheadLogTask.builder().config(config).streamRecordsMap(streamRecordsMap) + .objectManager(objectManager).objectStorage(objectStorage).executor(uploadWALExecutor).rate(rate).build(); + } + @SuppressWarnings("UnusedReturnValue") CompletableFuture uploadDeltaWAL() { return uploadDeltaWAL(LogCache.MATCH_ALL_STREAMS, false); } CompletableFuture uploadDeltaWAL(long streamId, boolean force) { + CompletableFuture cf; + List lazyUploadTasks = new ArrayList<>(); + lazyUploadQueue.drainTo(lazyUploadTasks); + synchronized (deltaWALCache) { - deltaWALCache.setConfirmOffset(confirmOffsetCalculator.get()); Optional blockOpt = deltaWALCache.archiveCurrentBlockIfContains(streamId); if (blockOpt.isPresent()) { LogCache.LogCacheBlock logCacheBlock = blockOpt.get(); DeltaWALUploadTaskContext context = new DeltaWALUploadTaskContext(logCacheBlock); context.objectManager = this.objectManager; context.force = force; - return uploadDeltaWAL(context); + cf = uploadDeltaWAL(context); } else { - return CompletableFuture.completedFuture(null); + cf = CompletableFuture.completedFuture(null); } } + + // notify lazy upload tasks + notifyLazyUpload(lazyUploadTasks); + return cf; } // only for test @@ -727,20 +948,13 @@ private void uploadDeltaWAL0(DeltaWALUploadTaskContext context) { if (context.force || elapsed <= 100L) { rate = Long.MAX_VALUE; } else { - rate = context.cache.size() * 1000.0 / Math.min(5000L, elapsed); + rate = context.cache.size() * 1000.0 / Math.min(20000L, elapsed); if (rate > maxDataWriteRate) { maxDataWriteRate = rate; } rate = maxDataWriteRate; } - context.task = DeltaWALUploadTask.builder() - .config(config) - .streamRecordsMap(context.cache.records()) - .objectManager(objectManager) - .objectStorage(objectStorage) - .executor(uploadWALExecutor) - .rate(rate) - .build(); + context.task = newUploadWriteAheadLogTask(context.cache.records(), objectManager, rate); boolean walObjectPrepareQueueEmpty = walPrepareQueue.isEmpty(); walPrepareQueue.add(context); if (!walObjectPrepareQueueEmpty) { @@ -779,9 +993,8 @@ private void commitDeltaWALUpload(DeltaWALUploadTaskContext context) { StorageOperationStats.getInstance().uploadWALCommitStats.record(context.timer.elapsedAs(TimeUnit.NANOSECONDS)); // 1. poll out current task walCommitQueue.poll(); - if (context.cache.confirmOffset() != 0) { - LOGGER.info("try trim WAL to {}", context.cache.confirmOffset()); - deltaWAL.trim(context.cache.confirmOffset()); + if (context.cache.lastRecordOffset() != null) { + delayTrim.trim(context.cache.lastRecordOffset(), context.trimCf); } // transfer records ownership to block cache. freeCache(context.cache); @@ -807,228 +1020,57 @@ private void freeCache(LogCache.LogCacheBlock cacheBlock) { deltaWALCache.markFree(cacheBlock); } - /** - * WALConfirmOffsetCalculator is used to calculate the confirmed offset of WAL. - */ - static class WALConfirmOffsetCalculator { - public static final long NOOP_OFFSET = -1L; - private final ReadWriteLock rwLock = new ReentrantReadWriteLock(); - private final Queue queue = new ConcurrentLinkedQueue<>(); - private final AtomicLong confirmOffset = new AtomicLong(NOOP_OFFSET); - - public WALConfirmOffsetCalculator() { - // Update the confirmed offset periodically. - Threads.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("wal-calculator-update-confirm-offset", true), LOGGER) - .scheduleAtFixedRate(this::update, 100, 100, TimeUnit.MILLISECONDS); - } - - /** - * Lock of {@link #add}. - * Operations of assigning offsets, for example {@link WriteAheadLog#append}, need to be performed while holding the lock. - */ - public Lock addLock() { - return rwLock.readLock(); - } + class DelayTrim { + private final long delayMillis; + private final BlockingQueue>> offsets = new LinkedBlockingQueue<>(); - public void add(WalWriteRequest request) { - assert null != request; - queue.add(new WalWriteRequestWrapper(request)); + public DelayTrim(long delayMillis) { + this.delayMillis = delayMillis; } - /** - * Return the offset before and including which all records have been persisted. - * Note: It is updated by {@link #update} periodically, and is not real-time. - */ - public Long get() { - return confirmOffset.get(); - } - - /** - * Calculate and update the confirmed offset. - */ - public void update() { - long offset = calculate(); - if (offset != NOOP_OFFSET) { - confirmOffset.set(offset); + public void trim(RecordOffset recordOffset, CompletableFuture cf) { + if (delayMillis == 0) { + LOGGER.info("try trim WAL to {}", recordOffset); + FutureUtil.propagate(deltaWAL.trim(recordOffset), cf); + } else { + offsets.add(Pair.of(recordOffset, cf)); + Threads.COMMON_SCHEDULER.schedule(() -> { + run(); + }, delayMillis, TimeUnit.MILLISECONDS); } } - /** - * Calculate the offset before and including which all records have been persisted. - * All records whose offset is not larger than the returned offset will be removed from the queue. - * It returns {@link #NOOP_OFFSET} if the first record is not persisted yet. - */ - private synchronized long calculate() { - Lock lock = rwLock.writeLock(); - lock.lock(); - try { - // Insert a flag. - queue.add(WalWriteRequestWrapper.flag()); - } finally { - lock.unlock(); - } - - long minUnconfirmedOffset = Long.MAX_VALUE; - boolean reachFlag = false; - for (WalWriteRequestWrapper wrapper : queue) { - // Iterate the queue to find the min unconfirmed offset. - if (wrapper.isFlag()) { - // Reach the flag. - reachFlag = true; - break; - } - WalWriteRequest request = wrapper.request; - assert request.offset != NOOP_OFFSET; - if (!request.confirmed) { - minUnconfirmedOffset = Math.min(minUnconfirmedOffset, request.offset); - } - } - assert reachFlag; - - long confirmedOffset = NOOP_OFFSET; - // Iterate the queue to find the max offset less than minUnconfirmedOffset. - // Remove all records whose offset is less than minUnconfirmedOffset. - for (Iterator iterator = queue.iterator(); iterator.hasNext(); ) { - WalWriteRequestWrapper wrapper = iterator.next(); - if (wrapper.isFlag()) { - /// Reach and remove the flag. - iterator.remove(); - break; - } - WalWriteRequest request = wrapper.request; - if (request.confirmed && request.offset < minUnconfirmedOffset) { - confirmedOffset = Math.max(confirmedOffset, request.offset); - iterator.remove(); - } + private void run() { + Pair> pair = offsets.poll(); + if (pair == null) { + return; } - return confirmedOffset; + LOGGER.info("try trim WAL to {}", pair.getKey()); + FutureUtil.propagate(deltaWAL.trim(pair.getKey()), pair.getValue()); } - /** - * Wrapper of {@link WalWriteRequest}. - * When the {@code request} is null, it is used as a flag. - */ - static final class WalWriteRequestWrapper { - private final WalWriteRequest request; - - /** - * - */ - WalWriteRequestWrapper(WalWriteRequest request) { - this.request = request; - } - - static WalWriteRequestWrapper flag() { - return new WalWriteRequestWrapper(null); - } - - public boolean isFlag() { - return request == null; - } - - public WalWriteRequest request() { - return request; - } - - @Override - public boolean equals(Object obj) { - if (obj == this) - return true; - if (obj == null || obj.getClass() != this.getClass()) - return false; - var that = (WalWriteRequestWrapper) obj; - return Objects.equals(this.request, that.request); - } - - @Override - public int hashCode() { - return Objects.hash(request); - } - - @Override - public String toString() { - return "WalWriteRequestWrapper[" + - "request=" + request + ']'; + public void close() { + List> futures = new ArrayList<>(); + List>> pending = new ArrayList<>(); + offsets.drainTo(pending); + for (Pair> pair : pending) { + FutureUtil.propagate(deltaWAL.trim(pair.getKey()), pair.getValue()); + futures.add(pair.getValue()); } - + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); } } - /** - * WALCallbackSequencer is used to sequence the unordered returned persistent data. - */ - static class WALCallbackSequencer { - private final Map> stream2requests = new ConcurrentHashMap<>(); - - /** - * Add request to stream sequence queue. - * When the {@code request.record.getStreamId()} is different, concurrent calls are allowed. - * When the {@code request.record.getStreamId()} is the same, concurrent calls are not allowed. And it is - * necessary to ensure that calls are made in the order of increasing offsets. - */ - public void before(WalWriteRequest request) { - try { - Queue streamRequests = stream2requests.computeIfAbsent(request.record.getStreamId(), - s -> new ConcurrentLinkedQueue<>()); - streamRequests.add(request); - } catch (Throwable ex) { - request.cf.completeExceptionally(ex); - } - } - - /** - * Try pop sequence persisted request from stream queue and move forward wal inclusive confirm offset. - * When the {@code request.record.getStreamId()} is different, concurrent calls are allowed. - * When the {@code request.record.getStreamId()} is the same, concurrent calls are not allowed. - * - * @return popped sequence persisted request. - */ - public List after(WalWriteRequest request) { - request.persisted = true; - - // Try to pop sequential persisted requests from the queue. - long streamId = request.record.getStreamId(); - Queue streamRequests = stream2requests.get(streamId); - WalWriteRequest peek = streamRequests.peek(); - if (peek == null || peek.offset != request.offset) { - return Collections.emptyList(); - } - - LinkedList rst = new LinkedList<>(); - WalWriteRequest poll = streamRequests.poll(); - assert poll == peek; - rst.add(poll); - - for (; ; ) { - peek = streamRequests.peek(); - if (peek == null || !peek.persisted) { - break; - } - poll = streamRequests.poll(); - assert poll == peek; - assert poll.record.getBaseOffset() == rst.getLast().record.getLastOffset(); - rst.add(poll); - } - - return rst; - } - - /** - * Try free stream related resources. - */ - public void tryFree(long streamId) { - Queue queue = stream2requests.get(streamId); - if (queue != null && queue.isEmpty()) { - stream2requests.remove(streamId, queue); - } - } + static StreamRecordBatch toLinkRecord(StreamRecordBatch origin, ByteBuf link) { + return StreamRecordBatch.of(origin.getStreamId(), origin.getEpoch(), origin.getBaseOffset(), -origin.getCount(), link); } public static class DeltaWALUploadTaskContext { TimerUtil timer; LogCache.LogCacheBlock cache; - DeltaWALUploadTask task; + UploadWriteAheadLogTask task; CompletableFuture cf; + CompletableFuture trimCf = new CompletableFuture<>(); ObjectManager objectManager; /** * Indicate whether to force upload the delta wal. @@ -1042,23 +1084,34 @@ public DeltaWALUploadTaskContext(LogCache.LogCacheBlock cache) { } /** - * Recover result of {@link #recoverContinuousRecords(Iterator, List, Logger)} - * Only streams not in {@link #invalidStreams} should be uploaded and closed. + * Recover result of {@link #recoverContinuousRecords(Iterator, Map, long, Logger)} */ - static class InnerRecoverResult { + static class RecoveryBlockResult { /** - * Recovered records. All {@link #invalidStreams} have been filtered out. + * Recovered records. All invalid streams have been filtered out. */ - LogCache.LogCacheBlock cacheBlock; + final LogCache.LogCacheBlock cacheBlock; /** - * Invalid streams, for example, the recovered start offset mismatches the stream end offset from controller. - * Key is streamId, value is the exception. + * Any exception occurred during recovery. It is null if no exception occurred. */ - Map invalidStreams = new HashMap<>(); + final RuntimeException exception; + + public RecoveryBlockResult(LogCache.LogCacheBlock cacheBlock, RuntimeException exception) { + this.cacheBlock = cacheBlock; + this.exception = exception; + } + } + + public static class LazyCommit { + final CompletableFuture trimCf = new CompletableFuture<>(); + final CompletableFuture commitCf = new CompletableFuture<>(); + final long lazyLingerMs; + final boolean awaitTrim; - public Optional firstException() { - return invalidStreams.values().stream().findFirst(); + public LazyCommit(long lazyLingerMs, boolean awaitTrim) { + this.lazyLingerMs = lazyLingerMs; + this.awaitTrim = awaitTrim; } } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/S3Stream.java b/s3stream/src/main/java/com/automq/stream/s3/S3Stream.java index 1434b50bb0..0a805782af 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/S3Stream.java +++ b/s3stream/src/main/java/com/automq/stream/s3/S3Stream.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -15,6 +23,7 @@ import com.automq.stream.RecordBatchWithContextWrapper; import com.automq.stream.api.AppendResult; import com.automq.stream.api.FetchResult; +import com.automq.stream.api.OpenStreamOptions; import com.automq.stream.api.RecordBatch; import com.automq.stream.api.RecordBatchWithContext; import com.automq.stream.api.Stream; @@ -24,20 +33,19 @@ import com.automq.stream.s3.cache.CacheAccessType; import com.automq.stream.s3.context.AppendContext; import com.automq.stream.s3.context.FetchContext; +import com.automq.stream.s3.metadata.StreamMetadata; import com.automq.stream.s3.metrics.S3StreamMetricsManager; import com.automq.stream.s3.metrics.TimerUtil; import com.automq.stream.s3.metrics.stats.NetworkStats; import com.automq.stream.s3.metrics.stats.StreamOperationStats; import com.automq.stream.s3.model.StreamRecordBatch; -import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; -import com.automq.stream.s3.network.NetworkBandwidthLimiter; -import com.automq.stream.s3.network.ThrottleStrategy; import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.s3.streams.StreamMetadataListener; import com.automq.stream.utils.FutureUtil; import com.automq.stream.utils.GlobalSwitch; +import com.automq.stream.utils.LogContext; import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import java.nio.ByteBuffer; import java.util.ArrayList; @@ -63,8 +71,8 @@ import static com.automq.stream.utils.FutureUtil.exec; import static com.automq.stream.utils.FutureUtil.propagate; -public class S3Stream implements Stream { - private static final Logger LOGGER = LoggerFactory.getLogger(S3Stream.class); +public class S3Stream implements Stream, StreamMetadataListener { + private final Logger logger; final AtomicLong confirmOffset; private final String logIdent; private final long streamId; @@ -79,33 +87,46 @@ public class S3Stream implements Stream { private final ReentrantLock appendLock = new ReentrantLock(); private final Set> pendingAppends = ConcurrentHashMap.newKeySet(); private final Deque pendingAppendTimestamps = new ConcurrentLinkedDeque<>(); + private volatile CompletableFuture lastAppendFuture; private final Set> pendingFetches = ConcurrentHashMap.newKeySet(); private final Deque pendingFetchTimestamps = new ConcurrentLinkedDeque<>(); - private final NetworkBandwidthLimiter networkInboundLimiter; - private final NetworkBandwidthLimiter networkOutboundLimiter; + private final OpenStreamOptions options; private long startOffset; private CompletableFuture lastPendingTrim = CompletableFuture.completedFuture(null); private CompletableFuture closeCf; + private StreamMetadataListener.Handle listenerHandle; - public S3Stream(long streamId, long epoch, long startOffset, long nextOffset, Storage storage, - StreamManager streamManager) { - this(streamId, epoch, startOffset, nextOffset, storage, streamManager, null, null); - } - - public S3Stream(long streamId, long epoch, long startOffset, long nextOffset, Storage storage, - StreamManager streamManager, NetworkBandwidthLimiter networkInboundLimiter, - NetworkBandwidthLimiter networkOutboundLimiter) { + private S3Stream(long streamId, long epoch, long startOffset, long nextOffset, Storage storage, + StreamManager streamManager, OpenStreamOptions options) { this.streamId = streamId; this.epoch = epoch; this.startOffset = startOffset; - this.logIdent = "[streamId=" + streamId + " epoch=" + epoch + "]"; + this.logIdent = "[streamId=" + streamId + " epoch=" + epoch + "] "; + this.logger = new LogContext(logIdent).logger(S3Stream.class); this.nextOffset = new AtomicLong(nextOffset); this.confirmOffset = new AtomicLong(nextOffset); this.status = new Status(); this.storage = storage; this.streamManager = streamManager; - this.networkInboundLimiter = networkInboundLimiter; - this.networkOutboundLimiter = networkOutboundLimiter; + this.options = options; + } + + public static S3Stream create(long streamId, long epoch, long startOffset, long nextOffset, Storage storage, + StreamManager streamManager) { + return create(streamId, epoch, startOffset, nextOffset, storage, streamManager, OpenStreamOptions.DEFAULT); + } + + public static S3Stream create(long streamId, long epoch, long startOffset, long nextOffset, Storage storage, + StreamManager streamManager, OpenStreamOptions options) { + S3Stream s3Stream = new S3Stream(streamId, epoch, startOffset, nextOffset, storage, streamManager, options); + s3Stream.completeInitialization(); + return s3Stream; + } + + private void completeInitialization() { + if (snapshotRead()) { + listenerHandle = streamManager.addMetadataListener(streamId, this); + } S3StreamMetricsManager.registerPendingStreamAppendLatencySupplier(streamId, () -> getHeadLatency(this.pendingAppendTimestamps)); S3StreamMetricsManager.registerPendingStreamFetchLatencySupplier(streamId, () -> getHeadLatency(this.pendingFetchTimestamps)); NetworkStats.getInstance().createStreamReadBytesStats(streamId); @@ -143,6 +164,14 @@ public long confirmOffset() { return this.confirmOffset.get(); } + @Override + public void confirmOffset(long offset) { + if (!snapshotRead()) { + throw new UnsupportedOperationException("Only snapshot-read mode support set confirmOffset"); + } + updateSnapshotReadConfirmOffset(offset); + } + @Override public long nextOffset() { return nextOffset.get(); @@ -151,20 +180,24 @@ public long nextOffset() { @Override @WithSpan public CompletableFuture append(AppendContext context, RecordBatch recordBatch) { + if (snapshotRead()) { + return FutureUtil.failedFuture(new IllegalStateException("Append operation is not support for readonly stream")); + } + if (recordBatch.count() < 0) { + return FutureUtil.failedFuture(new IllegalArgumentException("record batch count is negative")); + } long startTimeNanos = System.nanoTime(); readLock.lock(); try { CompletableFuture cf = exec(() -> { - if (networkInboundLimiter != null) { - networkInboundLimiter.consume(ThrottleStrategy.BYPASS, recordBatch.rawPayload().remaining()); - } appendLock.lock(); try { - return append0(context, recordBatch); + this.lastAppendFuture = append0(context, recordBatch); + return lastAppendFuture; } finally { appendLock.unlock(); } - }, LOGGER, "append"); + }, logger, "append"); pendingAppends.add(cf); pendingAppendTimestamps.push(startTimeNanos); return cf.whenComplete((nil, ex) -> { @@ -180,10 +213,10 @@ public CompletableFuture append(AppendContext context, RecordBatch @WithSpan private CompletableFuture append0(AppendContext context, RecordBatch recordBatch) { if (!status.isWritable()) { - return FutureUtil.failedFuture(new StreamClientException(ErrorCode.STREAM_ALREADY_CLOSED, logIdent + " stream is not writable")); + return FutureUtil.failedFuture(new StreamClientException(ErrorCode.STREAM_ALREADY_CLOSED, logIdent + "stream is not writable")); } long offset = nextOffset.getAndAdd(recordBatch.count()); - StreamRecordBatch streamRecordBatch = new StreamRecordBatch(streamId, epoch, offset, recordBatch.count(), Unpooled.wrappedBuffer(recordBatch.rawPayload())); + StreamRecordBatch streamRecordBatch = StreamRecordBatch.of(streamId, epoch, offset, recordBatch.count(), Unpooled.wrappedBuffer(recordBatch.rawPayload())); CompletableFuture cf = storage.append(context, streamRecordBatch).thenApply(nil -> { updateConfirmOffset(offset + recordBatch.count()); return new DefaultAppendResult(offset); @@ -196,45 +229,40 @@ private CompletableFuture append0(AppendContext context, RecordBat status.markFenced(); Throwable cause = FutureUtil.cause(ex); if (cause instanceof StreamClientException && ((StreamClientException) cause).getCode() == ErrorCode.EXPIRED_STREAM_EPOCH) { - LOGGER.info("{} stream append, stream is fenced", logIdent); + logger.info("stream append, stream is fenced"); } else { - LOGGER.warn("{} stream append fail", logIdent, cause); + logger.warn("stream append fail", cause); } }); } + @SuppressWarnings({"checkstyle:npathcomplexity"}) @Override @WithSpan public CompletableFuture fetch(FetchContext context, @SpanAttribute long startOffset, @SpanAttribute long endOffset, @SpanAttribute int maxBytes) { + if (snapshotRead()) { + context.readOptions().snapshotRead(true); + } TimerUtil timerUtil = new TimerUtil(); readLock.lock(); try { - CompletableFuture cf = exec(() -> fetch0(context, startOffset, endOffset, maxBytes), LOGGER, "fetch"); - CompletableFuture retCf = cf.thenCompose(rs -> { - if (networkOutboundLimiter != null) { - long totalSize = 0L; - for (RecordBatch recordBatch : rs.recordBatchList()) { - totalSize += recordBatch.rawPayload().remaining(); - } - final long finalSize = totalSize; - long start = System.nanoTime(); - ThrottleStrategy throttleStrategy = context.readOptions().prioritizedRead() ? ThrottleStrategy.BYPASS - : (context.readOptions().fastRead() ? ThrottleStrategy.TAIL : ThrottleStrategy.CATCH_UP); - return networkOutboundLimiter.consume(throttleStrategy, totalSize).thenApply(nil -> { - NetworkStats.getInstance().networkLimiterQueueTimeStats(AsyncNetworkBandwidthLimiter.Type.OUTBOUND, throttleStrategy) - .record(TimerUtil.timeElapsedSince(start, TimeUnit.NANOSECONDS)); - if (context.readOptions().fastRead()) { - NetworkStats.getInstance().fastReadBytesStats(streamId).ifPresent(counter -> counter.inc(finalSize)); - } else { - NetworkStats.getInstance().slowReadBytesStats(streamId).ifPresent(counter -> counter.inc(finalSize)); - } - return rs; - }); + CompletableFuture cf = exec(() -> fetch0(context, startOffset, endOffset, maxBytes), logger, "fetch"); + CompletableFuture retCf = cf.thenApply(rs -> { + // TODO: move the fast / slow read metrics to kafka module. + long totalSize = 0L; + for (RecordBatch recordBatch : rs.recordBatchList()) { + totalSize += recordBatch.rawPayload().remaining(); + } + final long finalSize = totalSize; + if (context.readOptions().fastRead()) { + NetworkStats.getInstance().fastReadBytesStats(streamId).ifPresent(counter -> counter.inc(finalSize)); + } else { + NetworkStats.getInstance().slowReadBytesStats(streamId).ifPresent(counter -> counter.inc(finalSize)); } - return CompletableFuture.completedFuture(rs); + return rs; }); pendingFetches.add(retCf); pendingFetchTimestamps.push(timerUtil.lastAs(TimeUnit.NANOSECONDS)); @@ -242,16 +270,16 @@ public CompletableFuture fetch(FetchContext context, if (ex != null) { Throwable cause = FutureUtil.cause(ex); if (!(cause instanceof FastReadFailFastException)) { - LOGGER.error("{} stream fetch [{}, {}) {} fail", logIdent, startOffset, endOffset, maxBytes, ex); + logger.error("stream fetch [{}, {}) {} fail", startOffset, endOffset, maxBytes, ex); } } StreamOperationStats.getInstance().fetchStreamLatency.record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - if (LOGGER.isDebugEnabled()) { + if (logger.isDebugEnabled()) { long totalSize = 0L; for (RecordBatch recordBatch : rs.recordBatchList()) { totalSize += recordBatch.rawPayload().remaining(); } - LOGGER.debug("fetch data, stream={}, {}-{}, total bytes: {}, cost={}ms", streamId, + logger.debug("fetch data, stream={}, {}-{}, total bytes: {}, cost={}ms", streamId, startOffset, endOffset, totalSize, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); } pendingFetches.remove(retCf); @@ -267,12 +295,12 @@ public CompletableFuture fetch(FetchContext context, private CompletableFuture fetch0(FetchContext context, long startOffset, long endOffset, int maxBytes) { if (!status.isReadable()) { - return FutureUtil.failedFuture(new StreamClientException(ErrorCode.STREAM_ALREADY_CLOSED, logIdent + " stream is already closed")); + return FutureUtil.failedFuture(new StreamClientException(ErrorCode.STREAM_ALREADY_CLOSED, logIdent + "stream is already closed")); } - if (LOGGER.isTraceEnabled()) { - LOGGER.trace("{} stream try fetch, startOffset: {}, endOffset: {}, maxBytes: {}", logIdent, startOffset, endOffset, maxBytes); + if (logger.isTraceEnabled()) { + logger.trace("stream try fetch, startOffset: {}, endOffset: {}, maxBytes: {}", startOffset, endOffset, maxBytes); } - long confirmOffset = this.confirmOffset.get(); + long confirmOffset = confirmOffset(); if (startOffset < startOffset() || endOffset > confirmOffset) { return FutureUtil.failedFuture( new StreamClientException( @@ -288,8 +316,8 @@ private CompletableFuture fetch0(FetchContext context, long startOf } return storage.read(context, streamId, startOffset, endOffset, maxBytes).thenApply(dataBlock -> { List records = dataBlock.getRecords(); - if (LOGGER.isTraceEnabled()) { - LOGGER.trace("{} stream fetch, startOffset: {}, endOffset: {}, maxBytes: {}, records: {}", logIdent, startOffset, endOffset, maxBytes, records.size()); + if (logger.isTraceEnabled()) { + logger.trace("stream fetch, startOffset: {}, endOffset: {}, maxBytes: {}, records: {}", startOffset, endOffset, maxBytes, records.size()); } return new DefaultFetchResult(records, dataBlock.getCacheAccessType(), context.readOptions().pooledBuf()); }); @@ -297,6 +325,9 @@ private CompletableFuture fetch0(FetchContext context, long startOf @Override public CompletableFuture trim(long newStartOffset) { + if (snapshotRead()) { + return FutureUtil.failedFuture(new IllegalStateException("Trim operation is not support for readonly stream")); + } writeLock.lock(); try { TimerUtil timerUtil = new TimerUtil(); @@ -306,7 +337,7 @@ public CompletableFuture trim(long newStartOffset) { this.lastPendingTrim = cf; cf.whenComplete((nil, ex) -> StreamOperationStats.getInstance().trimStreamLatency.record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS))); return cf; - }, LOGGER, "trim"); + }, logger, "trim"); } finally { writeLock.unlock(); } @@ -314,7 +345,7 @@ public CompletableFuture trim(long newStartOffset) { private CompletableFuture trim0(long newStartOffset) { if (newStartOffset < this.startOffset) { - LOGGER.warn("{} trim newStartOffset[{}] less than current start offset[{}]", logIdent, newStartOffset, startOffset); + logger.warn("trim newStartOffset[{}] less than current start offset[{}]", newStartOffset, startOffset); return CompletableFuture.completedFuture(null); } this.startOffset = newStartOffset; @@ -324,10 +355,10 @@ private CompletableFuture trim0(long newStartOffset) { awaitPendingFetchesCf.whenComplete((nil, ex) -> propagate(streamManager.trimStream(streamId, epoch, newStartOffset), trimCf)); trimCf.whenComplete((nil, ex) -> { if (ex != null) { - LOGGER.error("{} trim fail", logIdent, ex); + logger.error("trim fail", ex); } else { - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("{} trim to {}", logIdent, newStartOffset); + if (logger.isDebugEnabled()) { + logger.debug("trim to {}", newStartOffset); } } }); @@ -340,6 +371,13 @@ public CompletableFuture close() { } public CompletableFuture close(boolean force) { + if (snapshotRead()) { + listenerHandle.close(); + NetworkStats.getInstance().removeStreamReadBytesStats(streamId); + S3StreamMetricsManager.removePendingStreamAppendLatencySupplier(streamId); + S3StreamMetricsManager.removePendingStreamFetchLatencySupplier(streamId); + return CompletableFuture.completedFuture(null); + } TimerUtil timerUtil = new TimerUtil(); writeLock.lock(); try { @@ -347,10 +385,11 @@ public CompletableFuture close(boolean force) { // await all pending append/fetch/trim request List> pendingRequests = new ArrayList<>(pendingAppends); + // add timeout to prevent the fetch(catch-up read) network throttle to block Stream#close. if (GlobalSwitch.STRICT) { - pendingRequests.addAll(pendingFetches); + pendingRequests.addAll(FutureUtil.timeoutAndSilence(pendingFetches.stream(), 10, TimeUnit.SECONDS)); } - pendingRequests.add(lastPendingTrim); + pendingRequests.add(FutureUtil.timeoutAndSilence(lastPendingTrim, 10, TimeUnit.SECONDS)); if (force) { pendingRequests.forEach(cf -> cf.completeExceptionally(new StreamClientException(ErrorCode.UNEXPECTED, "FORCE_CLOSE"))); } @@ -366,14 +405,14 @@ public CompletableFuture close(boolean force) { CompletableFuture awaitPendingRequestsCf = CompletableFuture.allOf(pendingRequests.toArray(new CompletableFuture[0])); CompletableFuture closeCf = new CompletableFuture<>(); - awaitPendingRequestsCf.whenComplete((nil, ex) -> propagate(exec(this::close0, LOGGER, "close"), closeCf)); + awaitPendingRequestsCf.whenComplete((nil, ex) -> propagate(exec(this::close0, logger, "close"), closeCf)); closeCf.whenComplete((nil, ex) -> { if (ex != null) { - LOGGER.error("{} close fail", logIdent, ex); + logger.error("close fail", ex); StreamOperationStats.getInstance().closeStreamStats(false).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); } else { - LOGGER.info("{} closed", logIdent); + logger.info("closed"); StreamOperationStats.getInstance().closeStreamStats(true).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); } NetworkStats.getInstance().removeStreamReadBytesStats(streamId); @@ -395,14 +434,17 @@ private CompletableFuture close0() { @Override public CompletableFuture destroy() { + if (snapshotRead()) { + return FutureUtil.failedFuture(new IllegalStateException("Destroy operation is not support for readonly stream")); + } writeLock.lock(); try { - CompletableFuture destroyCf = close().thenCompose(nil -> exec(this::destroy0, LOGGER, "destroy")); + CompletableFuture destroyCf = close().thenCompose(nil -> exec(this::destroy0, logger, "destroy")); destroyCf.whenComplete((nil, ex) -> { if (ex != null) { - LOGGER.error("{} destroy fail", logIdent, ex); + logger.error("destroy fail", ex); } else { - LOGGER.info("{} destroyed", logIdent); + logger.info("destroyed"); } }); return destroyCf; @@ -411,6 +453,15 @@ public CompletableFuture destroy() { } } + @Override + public CompletableFuture lastAppendFuture() { + return lastAppendFuture; + } + + public boolean snapshotRead() { + return options.readWriteMode() == OpenStreamOptions.ReadWriteMode.SNAPSHOT_READ; + } + private CompletableFuture destroy0() { status.markDestroy(); startOffset = this.confirmOffset.get(); @@ -424,14 +475,27 @@ private void updateConfirmOffset(long newOffset) { break; } if (confirmOffset.compareAndSet(oldConfirmOffset, newOffset)) { - if (LOGGER.isTraceEnabled()) { - LOGGER.trace("{} stream update confirm offset from {} to {}", logIdent, oldConfirmOffset, newOffset); + if (logger.isTraceEnabled()) { + logger.trace("stream update confirm offset from {} to {}", oldConfirmOffset, newOffset); } break; } } } + @Override + public void onNewStreamMetadata(StreamMetadata metadata) { + updateSnapshotReadConfirmOffset(metadata.endOffset()); + this.startOffset = metadata.startOffset(); + } + + private void updateSnapshotReadConfirmOffset(long newOffset) { + synchronized (this.confirmOffset) { + this.confirmOffset.updateAndGet(operand -> Math.max(newOffset, operand)); + this.nextOffset.set(this.confirmOffset.get()); + } + } + static class DefaultFetchResult implements FetchResult { private static final LongAdder INFLIGHT = new LongAdder(); private final List pooledRecords; @@ -459,6 +523,8 @@ public DefaultFetchResult(List streamRecords, CacheAccessType private static RecordBatch convert(StreamRecordBatch streamRecordBatch, boolean pooledBuf) { ByteBuffer buf; + // We shouldn't access the StreamRecordBatch after release it. + int count = streamRecordBatch.getCount(); if (pooledBuf) { buf = streamRecordBatch.getPayload().nioBuffer(); } else { @@ -469,12 +535,12 @@ private static RecordBatch convert(StreamRecordBatch streamRecordBatch, boolean return new RecordBatch() { @Override public int count() { - return streamRecordBatch.getCount(); + return count; } @Override public long baseTimestamp() { - return streamRecordBatch.getEpoch(); + return 0; } @Override diff --git a/s3stream/src/main/java/com/automq/stream/s3/S3StreamClient.java b/s3stream/src/main/java/com/automq/stream/s3/S3StreamClient.java index ffa7afbeda..a0d7bc1a72 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/S3StreamClient.java +++ b/s3stream/src/main/java/com/automq/stream/s3/S3StreamClient.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -22,6 +30,7 @@ import com.automq.stream.s3.context.AppendContext; import com.automq.stream.s3.context.FetchContext; import com.automq.stream.s3.metadata.StreamMetadata; +import com.automq.stream.s3.metadata.StreamState; import com.automq.stream.s3.metrics.TimerUtil; import com.automq.stream.s3.metrics.stats.StreamOperationStats; import com.automq.stream.s3.network.NetworkBandwidthLimiter; @@ -109,6 +118,7 @@ public S3StreamClient(StreamManager streamManager, Storage storage, ObjectManage this.config = config; this.networkInboundBucket = networkInboundBucket; this.networkOutboundBucket = networkOutboundBucket; + ConfigValidator.validate(config); startStreamObjectsCompactions(); } @@ -119,7 +129,7 @@ public CompletableFuture createAndOpenStream(CreateStreamOptions options TimerUtil timerUtil = new TimerUtil(); return FutureUtil.exec(() -> streamManager.createStream(options.tags()).thenCompose(streamId -> { StreamOperationStats.getInstance().createStreamLatency.record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - return openStream0(streamId, options.epoch(), options.tags()); + return openStream0(streamId, options.epoch(), options.tags(), OpenStreamOptions.builder().epoch(options.epoch()).tags(options.tags()).build()); }), LOGGER, "createAndOpenStream"); }); } @@ -128,7 +138,7 @@ public CompletableFuture createAndOpenStream(CreateStreamOptions options public CompletableFuture openStream(long streamId, OpenStreamOptions openStreamOptions) { return runInLock(() -> { checkState(); - return FutureUtil.exec(() -> openStream0(streamId, openStreamOptions.epoch(), openStreamOptions.tags()), LOGGER, "openStream"); + return FutureUtil.exec(() -> openStream0(streamId, openStreamOptions.epoch(), openStreamOptions.tags(), openStreamOptions), LOGGER, "openStream"); }); } @@ -161,27 +171,42 @@ private void startStreamObjectsCompactions() { }, compactionJitterDelay, 1, TimeUnit.MINUTES); } - private CompletableFuture openStream0(long streamId, long epoch, Map tags) { + private CompletableFuture openStream0(long streamId, long epoch, Map tags, + OpenStreamOptions options) { return runInLock(() -> { TimerUtil timerUtil = new TimerUtil(); - CompletableFuture cf = streamManager.openStream(streamId, epoch, tags). - thenApply(metadata -> { - StreamWrapper stream = new StreamWrapper(newStream(metadata)); + CompletableFuture openStreamCf; + boolean snapshotRead = options.readWriteMode() == OpenStreamOptions.ReadWriteMode.SNAPSHOT_READ; + if (snapshotRead) { + openStreamCf = CompletableFuture.completedFuture(new StreamMetadata(streamId, epoch, -1, -1, StreamState.OPENED)); + } else { + openStreamCf = streamManager.openStream(streamId, epoch, tags); + } + CompletableFuture cf = openStreamCf.thenApply(metadata -> { + StreamWrapper stream = new StreamWrapper(newStream(metadata, options)); + if (!snapshotRead) { runInLock(() -> openedStreams.put(streamId, stream)); - StreamOperationStats.getInstance().openStreamLatency.record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - return stream; - }); - openingStreams.put(streamId, cf); - cf.whenComplete((stream, ex) -> runInLock(() -> openingStreams.remove(streamId, cf))); + } + StreamOperationStats.getInstance().openStreamLatency.record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + return stream; + }); + if (!snapshotRead) { + openingStreams.put(streamId, cf); + } + cf.whenComplete((stream, ex) -> runInLock(() -> { + if (!snapshotRead) { + openingStreams.remove(streamId, cf); + } + })); return cf; }); } - S3Stream newStream(StreamMetadata metadata) { - return new S3Stream( + S3Stream newStream(StreamMetadata metadata, OpenStreamOptions options) { + return S3Stream.create( metadata.streamId(), metadata.epoch(), metadata.startOffset(), metadata.endOffset(), - storage, streamManager, networkInboundBucket, networkOutboundBucket); + storage, streamManager, options); } @Override @@ -300,6 +325,11 @@ public long confirmOffset() { return stream.confirmOffset(); } + @Override + public void confirmOffset(long offset) { + stream.confirmOffset(offset); + } + @Override public long nextOffset() { return stream.nextOffset(); @@ -358,12 +388,17 @@ public CompletableFuture destroy() { }); } + @Override + public CompletableFuture lastAppendFuture() { + return stream.lastAppendFuture(); + } + public boolean isClosed() { return stream.isClosed(); } - public void compact(CompactionHint hint) { - if (isClosed()) { + void compact(CompactionHint hint) { + if (isClosed() || stream.snapshotRead()) { // the compaction task may be taking a long time, // so we need to check if the stream is closed before starting the compaction. return; @@ -405,7 +440,7 @@ private void compactV1(CompactionHint hint, long now) { } } - public void compact(StreamObjectCompactor.CompactionType compactionType, CompactionHint hint) { + private void compact(StreamObjectCompactor.CompactionType compactionType, CompactionHint hint) { StreamObjectCompactor.Builder taskBuilder = StreamObjectCompactor.builder() .objectManager(objectManager) .stream(this) diff --git a/s3stream/src/main/java/com/automq/stream/s3/Storage.java b/s3stream/src/main/java/com/automq/stream/s3/Storage.java index e6bf585cd7..6a6c8376c0 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/Storage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/Storage.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/main/java/com/automq/stream/s3/StreamDataBlock.java b/s3stream/src/main/java/com/automq/stream/s3/StreamDataBlock.java index 00b8472fe0..2cec663af0 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/StreamDataBlock.java +++ b/s3stream/src/main/java/com/automq/stream/s3/StreamDataBlock.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -81,6 +89,13 @@ public CompletableFuture getDataCf() { return this.dataCf; } + public ByteBuf getAndReleaseData() { + if (refCount.getAndDecrement() == 0) { + throw new IllegalStateException("Data has already been released"); + } + return this.dataCf.join(); + } + public void releaseRef() { refCount.decrementAndGet(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/StreamRecordBatchCodec.java b/s3stream/src/main/java/com/automq/stream/s3/StreamRecordBatchCodec.java index 564de418d7..008499a377 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/StreamRecordBatchCodec.java +++ b/s3stream/src/main/java/com/automq/stream/s3/StreamRecordBatchCodec.java @@ -1,23 +1,24 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; -import com.automq.stream.ByteBufSeqAlloc; -import com.automq.stream.s3.model.StreamRecordBatch; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.ByteBufAlloc.ENCODE_RECORD; - public class StreamRecordBatchCodec { public static final byte MAGIC_V0 = 0x22; public static final int HEADER_SIZE = @@ -27,54 +28,11 @@ public class StreamRecordBatchCodec { + 8 // baseOffset + 4 // lastOffsetDelta + 4; // payload length - private static final ByteBufSeqAlloc ENCODE_ALLOC = new ByteBufSeqAlloc(ENCODE_RECORD, 8); - - public static ByteBuf encode(StreamRecordBatch streamRecord) { - int totalLength = HEADER_SIZE + streamRecord.size(); // payload - // use sequential allocator to avoid memory fragmentation - ByteBuf buf = ENCODE_ALLOC.byteBuffer(totalLength); - buf.writeByte(MAGIC_V0); - buf.writeLong(streamRecord.getStreamId()); - buf.writeLong(streamRecord.getEpoch()); - buf.writeLong(streamRecord.getBaseOffset()); - buf.writeInt(streamRecord.getCount()); - buf.writeInt(streamRecord.size()); - buf.writeBytes(streamRecord.getPayload().duplicate()); - return buf; - } - - /** - * Decode a stream record batch from a byte buffer and move the reader index. - * The returned stream record batch does NOT share the payload buffer with the input buffer. - */ - public static StreamRecordBatch duplicateDecode(ByteBuf buf) { - byte magic = buf.readByte(); // magic - if (magic != MAGIC_V0) { - throw new RuntimeException("Invalid magic byte " + magic); - } - long streamId = buf.readLong(); - long epoch = buf.readLong(); - long baseOffset = buf.readLong(); - int lastOffsetDelta = buf.readInt(); - int payloadLength = buf.readInt(); - ByteBuf payload = ByteBufAlloc.byteBuffer(payloadLength, ByteBufAlloc.DECODE_RECORD); - buf.readBytes(payload); - return new StreamRecordBatch(streamId, epoch, baseOffset, lastOffsetDelta, payload); - } - - /** - * Decode a stream record batch from a byte buffer and move the reader index. - * The returned stream record batch shares the payload buffer with the input buffer. - */ - public static StreamRecordBatch decode(ByteBuf buf) { - buf.readByte(); // magic - long streamId = buf.readLong(); - long epoch = buf.readLong(); - long baseOffset = buf.readLong(); - int lastOffsetDelta = buf.readInt(); - int payloadLength = buf.readInt(); - ByteBuf payload = buf.slice(buf.readerIndex(), payloadLength); - buf.skipBytes(payloadLength); - return new StreamRecordBatch(streamId, epoch, baseOffset, lastOffsetDelta, payload); - } + public static final int MAGIC_POS = 0; + public static final int STREAM_ID_POS = 1; + public static final int EPOCH_POS = STREAM_ID_POS + 8; + public static final int BASE_OFFSET_POS = EPOCH_POS + 8; + public static final int LAST_OFFSET_DELTA_POS = BASE_OFFSET_POS + 8; + public static final int PAYLOAD_LENGTH_POS = LAST_OFFSET_DELTA_POS + 4; + public static final int PAYLOAD_POS = PAYLOAD_LENGTH_POS + 4; } diff --git a/s3stream/src/main/java/com/automq/stream/s3/UploadWriteAheadLogTask.java b/s3stream/src/main/java/com/automq/stream/s3/UploadWriteAheadLogTask.java new file mode 100644 index 0000000000..684deea6d1 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/UploadWriteAheadLogTask.java @@ -0,0 +1,48 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3; + +import com.automq.stream.s3.objects.CommitStreamSetObjectRequest; + +import java.util.concurrent.CompletableFuture; + +public interface UploadWriteAheadLogTask { + /** + * Prepare the upload task and get the uploadId(objectId). + * @return the uploadId(objectId) + */ + CompletableFuture prepare(); + + /** + * Upload the delta wal data to the main storage + */ + CompletableFuture upload(); + + /** + * Commit the upload result to controller. + */ + CompletableFuture commit(); + + /** + * Burst the upload task. + */ + void burst(); + +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/WalWriteRequest.java b/s3stream/src/main/java/com/automq/stream/s3/WalWriteRequest.java index e2bd49bbaf..0fbeda7cec 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/WalWriteRequest.java +++ b/s3stream/src/main/java/com/automq/stream/s3/WalWriteRequest.java @@ -1,20 +1,27 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; -import com.automq.stream.s3.cache.LogCache; import com.automq.stream.s3.context.AppendContext; import com.automq.stream.s3.model.StreamRecordBatch; -import com.automq.stream.s3.wal.WriteAheadLog; +import com.automq.stream.s3.wal.RecordOffset; import java.util.concurrent.CompletableFuture; @@ -22,28 +29,9 @@ public class WalWriteRequest implements Comparable { final StreamRecordBatch record; final AppendContext context; final CompletableFuture cf; - long offset; - /** - * Whether the record has been persisted to the {@link WriteAheadLog} - * When a continuous series of records IN A STREAM have been persisted to the WAL, they can be uploaded to S3. - * - * @see S3Storage.WALCallbackSequencer - */ - boolean persisted; - - /** - * Whether the record has been put to the {@link LogCache} - * When a continuous series of records have been persisted to the WAL and uploaded to S3, they can be trimmed. - * - * @see S3Storage.WALConfirmOffsetCalculator - */ - boolean confirmed; - - public WalWriteRequest(StreamRecordBatch record, long offset, CompletableFuture cf) { - this(record, offset, cf, AppendContext.DEFAULT); - } + RecordOffset offset; - public WalWriteRequest(StreamRecordBatch record, long offset, CompletableFuture cf, AppendContext context) { + public WalWriteRequest(StreamRecordBatch record, RecordOffset offset, CompletableFuture cf, AppendContext context) { this.record = record; this.offset = offset; this.cf = cf; @@ -60,8 +48,6 @@ public String toString() { return "WalWriteRequest{" + "record=" + record + ", offset=" + offset + - ", persisted=" + persisted + - ", confirmed=" + confirmed + '}'; } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncLRUCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncLRUCache.java index fdf9cdbc2e..42b70f988e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncLRUCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncLRUCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; @@ -22,7 +30,14 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Function; +/** + * An asynchronous LRU cache that supports asynchronous value computation. + * + * @param key type + * @param value type, NOTE: V must not override equals and hashCode + */ public class AsyncLRUCache { private static final Logger LOGGER = LoggerFactory.getLogger(AsyncLRUCache.class); private final AsyncLRUCacheStats stats = AsyncLRUCacheStats.getInstance(); @@ -33,13 +48,21 @@ public class AsyncLRUCache { final Set completedSet = new HashSet<>(); final Set removedSet = new HashSet<>(); - public AsyncLRUCache(String cacheName, long maxSize) { + protected AsyncLRUCache(String cacheName, long maxSize) { this.cacheName = cacheName; if (maxSize <= 0) { throw new IllegalArgumentException("maxSize must be positive"); } this.maxSize = maxSize; + } + public static AsyncLRUCache create(String cacheName, long maxSize) { + AsyncLRUCache asyncLRUCache = new AsyncLRUCache<>(cacheName, maxSize); + asyncLRUCache.completeInitialization(); + return asyncLRUCache; + } + + private void completeInitialization() { S3StreamMetricsManager.registerAsyncCacheSizeSupplier(this::totalSize, cacheName); S3StreamMetricsManager.registerAsyncCacheMaxSizeSupplier(() -> maxSize, cacheName); S3StreamMetricsManager.registerAsyncCacheItemNumberSupplier(this::size, cacheName); @@ -88,6 +111,20 @@ public synchronized V get(K key) { return val; } + public synchronized V computeIfAbsent(K key, Function valueMapper) { + V value = cache.get(key); + if (value == null) { + value = valueMapper.apply(key); + if (value != null) { + put(key, value); + } + } + return value; + } + + public synchronized void inLockRun(Runnable runnable) { + runnable.run(); + } public synchronized boolean remove(K key) { V value = cache.get(key); diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncMeasurable.java b/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncMeasurable.java index 599ffa5f49..0ce16ff78e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncMeasurable.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/AsyncMeasurable.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/CacheAccessType.java b/s3stream/src/main/java/com/automq/stream/s3/cache/CacheAccessType.java index 834bb4d686..866fa0fe28 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/CacheAccessType.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/CacheAccessType.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/LRUCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/LRUCache.java index 1d1f7ab9b7..7c1db3d69b 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/LRUCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/LRUCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/LogCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/LogCache.java index 490430383a..672242b4a1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/LogCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/LogCache.java @@ -1,21 +1,30 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; -import com.automq.stream.s3.metrics.S3StreamMetricsManager; import com.automq.stream.s3.metrics.TimerUtil; import com.automq.stream.s3.metrics.stats.StorageOperationStats; import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.trace.context.TraceContext; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.utils.Threads; import com.automq.stream.utils.biniarysearch.StreamRecordBatchList; import org.slf4j.Logger; @@ -29,7 +38,9 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -49,7 +60,12 @@ public class LogCache { private static final int DEFAULT_MAX_BLOCK_STREAM_COUNT = 10000; private static final Consumer DEFAULT_BLOCK_FREE_LISTENER = block -> { }; + private static final int MAX_BLOCKS_COUNT = 64; + private static final ExecutorService LOG_CACHE_ASYNC_EXECUTOR = Threads.newFixedFastThreadLocalThreadPoolWithMonitor( + 1, "LOG_CACHE_ASYNC", true, LOGGER); + static final int MERGE_BLOCK_THRESHOLD = 8; final List blocks = new ArrayList<>(); + final AtomicInteger blockCount = new AtomicInteger(1); private final long capacity; private final long cacheBlockMaxSize; private final int maxCacheBlockStreamCount; @@ -60,7 +76,7 @@ public class LogCache { private final ReentrantReadWriteLock.ReadLock readLock = lock.readLock(); private final ReentrantReadWriteLock.WriteLock writeLock = lock.writeLock(); private LogCacheBlock activeBlock; - private long confirmOffset; + private RecordOffset lastRecordOffset; public LogCache(long capacity, long cacheBlockMaxSize, int maxCacheBlockStreamCount, Consumer blockFreeListener) { @@ -70,7 +86,6 @@ public LogCache(long capacity, long cacheBlockMaxSize, int maxCacheBlockStreamCo this.activeBlock = new LogCacheBlock(cacheBlockMaxSize, maxCacheBlockStreamCount); this.blocks.add(activeBlock); this.blockFreeListener = blockFreeListener; - S3StreamMetricsManager.registerDeltaWalCacheSizeSupplier(size::get); } public LogCache(long capacity, long cacheBlockMaxSize) { @@ -139,7 +154,8 @@ public List get(TraceContext context, List records; readLock.lock(); try { - records = get0(streamId, startOffset, endOffset, maxBytes); + Long streamIdLong = streamId; + records = get0(streamIdLong, startOffset, endOffset, maxBytes); records.forEach(StreamRecordBatch::retain); } finally { readLock.unlock(); @@ -151,7 +167,7 @@ public List get(TraceContext context, return records; } - public List get0(long streamId, long startOffset, long endOffset, int maxBytes) { + public List get0(Long streamId, long startOffset, long endOffset, int maxBytes) { List rst = new LinkedList<>(); long nextStartOffset = startOffset; int nextMaxBytes = maxBytes; @@ -203,9 +219,10 @@ public LogCacheBlock archiveCurrentBlock() { writeLock.lock(); try { LogCacheBlock block = activeBlock; - block.confirmOffset = confirmOffset; + block.lastRecordOffset = lastRecordOffset; activeBlock = new LogCacheBlock(cacheBlockMaxSize, maxCacheBlockStreamCount); blocks.add(activeBlock); + blockCount.set(blocks.size()); return block; } finally { writeLock.unlock(); @@ -238,77 +255,167 @@ Optional archiveCurrentBlockIfContains0(long streamId) { } - public void markFree(LogCacheBlock block) { + public CompletableFuture markFree(LogCacheBlock block) { block.free = true; tryRealFree(); + CompletableFuture cf = new CompletableFuture<>(); + LOG_CACHE_ASYNC_EXECUTOR.execute(() -> { + try { + tryMerge(); + cf.complete(null); + } catch (Throwable t) { + cf.completeExceptionally(t); + } + }); + return cf; } private void tryRealFree() { long currSize = size.get(); - if (currSize <= capacity * 0.9) { + if (currSize <= capacity * 0.9 && blockCount.get() <= MAX_BLOCKS_COUNT) { return; } List removed = new ArrayList<>(); long freeSize = 0L; writeLock.lock(); try { + // free blocks currSize = size.get(); Iterator iter = blocks.iterator(); while (iter.hasNext()) { - if (currSize - freeSize <= capacity * 0.9) { + LogCacheBlock block = iter.next(); + if (blockCount.get() <= MAX_BLOCKS_COUNT && currSize - freeSize <= capacity * 0.9) { break; } - LogCacheBlock block = iter.next(); if (block.free) { iter.remove(); freeSize += block.size(); removed.add(block); + blockCount.decrementAndGet(); } else { break; } } - } finally { writeLock.unlock(); } size.addAndGet(-freeSize); - removed.forEach(b -> { + LOG_CACHE_ASYNC_EXECUTOR.execute(() -> removed.forEach(b -> { blockFreeListener.accept(b); b.free(); - }); + })); } - public int forceFree(int required) { - AtomicInteger freedBytes = new AtomicInteger(); - List removed = new ArrayList<>(); + private void tryMerge() { + // merge blocks to speed up the get. + int mergeStartIndex = 0; + for (; ; ) { + LogCacheBlock left; + LogCacheBlock right; + writeLock.lock(); + try { + if (blocks.size() <= MERGE_BLOCK_THRESHOLD || mergeStartIndex + 1 >= blocks.size()) { + return; + } + left = blocks.get(mergeStartIndex); + right = blocks.get(mergeStartIndex + 1); + if (!left.free || !right.free) { + return; + } + if (left.size() + right.size() >= cacheBlockMaxSize) { + mergeStartIndex++; + continue; + } + } finally { + writeLock.unlock(); + } + // Move costly operation(isDiscontinuous, mergeBlock) out of the lock. + if (isDiscontinuous(left, right)) { + mergeStartIndex++; + continue; + } + LogCacheBlock newBlock = new LogCacheBlock(Integer.MAX_VALUE); + mergeBlock(newBlock, left); + mergeBlock(newBlock, right); + newBlock.free = true; + writeLock.lock(); + try { + if (blocks.size() > mergeStartIndex + 1 + && blocks.get(mergeStartIndex) == left + && blocks.get(mergeStartIndex + 1) == right + ) { + blocks.set(mergeStartIndex, newBlock); + blocks.remove(mergeStartIndex + 1); + } + } finally { + writeLock.unlock(); + } + } + } + + public void setLastRecordOffset(RecordOffset lastRecordOffset) { + readLock.lock(); + try { + this.lastRecordOffset = lastRecordOffset; + } finally { + readLock.unlock(); + } + } + + public long size() { + return size.get(); + } + + public long capacity() { + return capacity; + } + + public void clearStreamRecords(long streamId) { writeLock.lock(); try { - blocks.removeIf(block -> { - if (!block.free || freedBytes.get() >= required) { - return false; - } - long blockSize = block.size(); - size.addAndGet(-blockSize); - freedBytes.addAndGet((int) blockSize); - removed.add(block); - return true; - }); + for (LogCacheBlock block : blocks) { + block.free(streamId); + } } finally { writeLock.unlock(); } - removed.forEach(b -> { - blockFreeListener.accept(b); - b.free(); - }); - return freedBytes.get(); } - public void setConfirmOffset(long confirmOffset) { - this.confirmOffset = confirmOffset; + static boolean isDiscontinuous(LogCacheBlock left, LogCacheBlock right) { + for (Map.Entry entry : left.map.entrySet()) { + Long streamId = entry.getKey(); + StreamCache leftStreamCache = entry.getValue(); + StreamCache rightStreamCache = right.map.get(streamId); + if (rightStreamCache == null) { + continue; + } + if (leftStreamCache.endOffset() != rightStreamCache.startOffset()) { + return true; + } + } + return false; } - public long size() { - return size.get(); + /** + * Merge the right block to the left block + */ + static void mergeBlock(LogCacheBlock left, LogCacheBlock right) { + synchronized (left) { + left.size.addAndGet(right.size()); + left.lastRecordOffset = right.lastRecordOffset; + left.map.forEach((streamId, leftStreamCache) -> { + StreamCache rightStreamCache = right.map.get(streamId); + if (rightStreamCache != null) { + leftStreamCache.records.addAll(rightStreamCache.records); + leftStreamCache.endOffset(rightStreamCache.endOffset()); + } + }); + right.map.forEach((streamId, rightStreamCache) -> { + if (!left.map.containsKey(streamId)) { + left.map.put(streamId, rightStreamCache); + } + }); + } } public static class LogCacheBlock { @@ -319,8 +426,9 @@ public static class LogCacheBlock { private final int maxStreamCount; private final long createdTimestamp = System.currentTimeMillis(); private final AtomicLong size = new AtomicLong(); + private final List freeListeners = new ArrayList<>(); volatile boolean free; - private long confirmOffset; + private RecordOffset lastRecordOffset; public LogCacheBlock(long maxSize, int maxStreamCount) { this.blockId = BLOCK_ID_ALLOC.getAndIncrement(); @@ -336,6 +444,10 @@ public long blockId() { return blockId; } + public boolean isFull() { + return size.get() >= maxSize || map.size() >= maxStreamCount; + } + public boolean put(StreamRecordBatch recordBatch) { map.compute(recordBatch.getStreamId(), (id, cache) -> { if (cache == null) { @@ -344,10 +456,11 @@ public boolean put(StreamRecordBatch recordBatch) { cache.add(recordBatch); return cache; }); - return size.addAndGet(recordBatch.occupiedSize()) >= maxSize || map.size() >= maxStreamCount; + size.addAndGet(recordBatch.occupiedSize()); + return isFull(); } - public List get(long streamId, long startOffset, long endOffset, int maxBytes) { + public List get(Long streamId, long startOffset, long endOffset, int maxBytes) { StreamCache cache = map.get(streamId); if (cache == null) { return Collections.emptyList(); @@ -355,7 +468,7 @@ public List get(long streamId, long startOffset, long endOffs return cache.get(startOffset, endOffset, maxBytes); } - StreamRange getStreamRange(long streamId) { + StreamRange getStreamRange(Long streamId) { StreamCache streamCache = map.get(streamId); if (streamCache == null) { return new StreamRange(NOOP_OFFSET, NOOP_OFFSET); @@ -370,12 +483,12 @@ public Map> records() { .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); } - public long confirmOffset() { - return confirmOffset; + public RecordOffset lastRecordOffset() { + return lastRecordOffset; } - public void confirmOffset(long confirmOffset) { - this.confirmOffset = confirmOffset; + public void lastRecordOffset(RecordOffset lastRecordOffset) { + this.lastRecordOffset = lastRecordOffset; } public long size() { @@ -384,11 +497,27 @@ public long size() { public void free() { suppress(() -> { - map.forEach((streamId, records) -> records.free()); + List streams = new ArrayList<>(map.size()); + map.forEach((streamId, records) -> { + streams.add(new StreamRangeBound(streamId, records.startOffset(), records.endOffset())); + records.free(); + }); map.clear(); + freeListeners.forEach(listener -> listener.onFree(streams)); }, LOGGER); } + public void free(long streamId) { + StreamCache streamCache = map.remove(streamId); + if (streamCache != null) { + LOG_CACHE_ASYNC_EXECUTOR.execute(() -> suppress(streamCache::free, LOGGER)); + } + } + + public void addFreeListener(FreeListener freeListener) { + freeListeners.add(freeListener); + } + public long createdTimestamp() { return createdTimestamp; } @@ -401,8 +530,37 @@ public boolean containsStream(long streamId) { } } + public interface FreeListener { + void onFree(List streamRanges); + } + + public static class StreamRangeBound { + private final long streamId; + private final long startOffset; + private final long endOffset; + + public StreamRangeBound(long streamId, long startOffset, long endOffset) { + this.streamId = streamId; + this.startOffset = startOffset; + this.endOffset = endOffset; + } + + public long streamId() { + return streamId; + } + + public long startOffset() { + return startOffset; + } + + public long endOffset() { + return endOffset; + } + } + static class StreamRange { public static final long NOOP_OFFSET = -1L; + public static final StreamRange NOOP = new StreamRange(NOOP_OFFSET, NOOP_OFFSET); long startOffset; long endOffset; @@ -413,11 +571,15 @@ public StreamRange(long startOffset, long endOffset) { } static class StreamCache { - List records = new ArrayList<>(); + List records; long startOffset = NOOP_OFFSET; long endOffset = NOOP_OFFSET; Map offsetIndexMap = new HashMap<>(); + public StreamCache() { + this.records = new ArrayList<>(); + } + synchronized void add(StreamRecordBatch recordBatch) { if (recordBatch.getBaseOffset() != endOffset && endOffset != NOOP_OFFSET) { RuntimeException ex = new IllegalArgumentException(String.format("streamId=%s record batch base offset mismatch, expect %s, actual %s", @@ -495,6 +657,22 @@ synchronized void free() { records.forEach(StreamRecordBatch::release); records.clear(); } + + synchronized long startOffset() { + return startOffset; + } + + synchronized long endOffset() { + return endOffset; + } + + synchronized void endOffset(long endOffset) { + this.endOffset = endOffset; + } + + synchronized int count() { + return records.size(); + } } static class IndexAndCount { diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/ObjectReaderLRUCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/ObjectReaderLRUCache.java index c6e5b6f5ad..0d02652c1a 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/ObjectReaderLRUCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/ObjectReaderLRUCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/ReadDataBlock.java b/s3stream/src/main/java/com/automq/stream/s3/cache/ReadDataBlock.java index 33a760eb7e..2c1032ca4c 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/ReadDataBlock.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/ReadDataBlock.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java index 4715d64d82..8770b43e54 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/SnapshotReadCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/SnapshotReadCache.java new file mode 100644 index 0000000000..50e657b730 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/SnapshotReadCache.java @@ -0,0 +1,490 @@ +package com.automq.stream.s3.cache; + +import com.automq.stream.ByteBufSeqAlloc; +import com.automq.stream.api.LinkRecordDecoder; +import com.automq.stream.s3.DataBlockIndex; +import com.automq.stream.s3.ObjectReader; +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.s3.metadata.StreamMetadata; +import com.automq.stream.s3.metrics.Metrics; +import com.automq.stream.s3.metrics.MetricsLevel; +import com.automq.stream.s3.metrics.wrapper.DeltaHistogram; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.WriteAheadLog; +import com.automq.stream.utils.CloseableIterator; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Time; +import com.automq.stream.utils.threads.EventLoop; +import com.automq.stream.utils.threads.EventLoopSafe; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.RemovalListener; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedList; +import java.util.List; +import java.util.ListIterator; +import java.util.Map; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Function; +import java.util.stream.Collectors; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; + +import static com.automq.stream.s3.ByteBufAlloc.SNAPSHOT_READ_CACHE; + +public class SnapshotReadCache { + public static final ByteBufSeqAlloc ENCODE_ALLOC = new ByteBufSeqAlloc(SNAPSHOT_READ_CACHE, 8); + private static final Logger LOGGER = LoggerFactory.getLogger(SnapshotReadCache.class); + private static final long MAX_INFLIGHT_LOAD_BYTES = 100L * 1024 * 1024; + + private static final Metrics.HistogramBundle OPERATION_LATENCY = Metrics.instance().histogram("kafka_stream_snapshot_read_cache", "Snapshot read cache operation latency", "nanoseconds"); + private static final DeltaHistogram REPLAY_LATENCY = OPERATION_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "replay")); + private static final DeltaHistogram READ_WAL_LATENCY = OPERATION_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "read_wal")); + private static final DeltaHistogram DECODE_LATENCY = OPERATION_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "decode")); + private static final DeltaHistogram PUT_INTO_CACHE_LATENCY = OPERATION_LATENCY.histogram(MetricsLevel.INFO, Attributes.of(AttributeKey.stringKey("operation"), "put_into_cache")); + + private final Map streamNextOffsets = new HashMap<>(); + private final Cache activeStreams; + private final EventLoop eventLoop = new EventLoop("SNAPSHOT_READ_CACHE"); + private final LogCacheBlockFreeListener cacheFreeListener = new LogCacheBlockFreeListener(); + + private final ObjectReplay objectReplay = new ObjectReplay(); + private final WalReplay walReplay = new WalReplay(); + private final StreamManager streamManager; + private final LogCache cache; + private final ObjectStorage objectStorage; + private final LinkRecordDecoder linkRecordDecoder; + private final Time time = Time.SYSTEM; + + public SnapshotReadCache(StreamManager streamManager, LogCache cache, ObjectStorage objectStorage, + LinkRecordDecoder linkRecordDecoder) { + activeStreams = CacheBuilder.newBuilder() + .expireAfterAccess(10, TimeUnit.MINUTES) + .removalListener((RemovalListener) notification -> + eventLoop.execute(() -> clearStream(notification.getKey()))) + .build(); + this.streamManager = streamManager; + this.cache = cache; + this.objectStorage = objectStorage; + this.linkRecordDecoder = linkRecordDecoder; + } + + @EventLoopSafe + void put(CloseableIterator it) { + long startNanos = time.nanoseconds(); + try (it) { + LogCache cache = this.cache; + if (cache == null) { + it.forEachRemaining(StreamRecordBatch::release); + return; + } + long streamId = -1L; + AtomicLong expectedNextOffset = null; + while (it.hasNext()) { + StreamRecordBatch batch = it.next(); + long newStreamId = batch.getStreamId(); + if (streamId == -1L || newStreamId != streamId) { + streamId = newStreamId; + expectedNextOffset = streamNextOffsets.computeIfAbsent(streamId, k -> new AtomicLong(batch.getBaseOffset())); + activeStream(streamId); + } + if (batch.getBaseOffset() < expectedNextOffset.get()) { + batch.release(); + continue; + } else if (batch.getBaseOffset() > expectedNextOffset.get()) { + // The LogCacheBlock doesn't accept discontinuous record batches. + cache.clearStreamRecords(streamId); + } + if (cache.put(batch)) { + // the block is full + LogCache.LogCacheBlock cacheBlock = cache.archiveCurrentBlock(); + cacheBlock.addFreeListener(cacheFreeListener); + cache.markFree(cacheBlock); + } + expectedNextOffset.set(batch.getLastOffset()); + } + } + PUT_INTO_CACHE_LATENCY.record(time.nanoseconds() - startNanos); + } + + public synchronized CompletableFuture replay(List objects) { + return objectReplay.replay(objects); + } + + public synchronized CompletableFuture replay(WriteAheadLog confirmWAL, RecordOffset startOffset, + RecordOffset endOffset, List walRecords) { + long startNanos = time.nanoseconds(); + return walReplay.replay(confirmWAL, startOffset, endOffset, walRecords) + .whenComplete((nil, ex) -> REPLAY_LATENCY.record(time.nanoseconds() - startNanos)); + } + + public void addEventListener(EventListener eventListener) { + cacheFreeListener.addListener(eventListener); + } + + @EventLoopSafe + private void clearStream(Long streamId) { + cache.clearStreamRecords(streamId); + streamNextOffsets.remove(streamId); + } + + private void activeStream(long streamId) { + try { + activeStreams.get(streamId, () -> true); + } catch (ExecutionException e) { + // suppress + } + } + + class WalReplay { + private static final long TASK_WAITING_TIMEOUT_NANOS = TimeUnit.SECONDS.toNanos(5); + private static final int MAX_WAITING_LOAD_TASK_COUNT = 4096; + // soft limit the inflight memory + private final int maxInflightLoadingCount = Systems.CPU_CORES * 4; + private final BlockingQueue waitingLoadTasks = new ArrayBlockingQueue<>(MAX_WAITING_LOAD_TASK_COUNT); + private final Queue loadingTasks = new ConcurrentLinkedQueue<>(); + + public CompletableFuture replay(WriteAheadLog wal, RecordOffset startOffset, RecordOffset endOffset, + List walRecords) { + WalReplayTask task = new WalReplayTask(wal, startOffset, endOffset, walRecords); + while (!waitingLoadTasks.add(task)) { + // The replay won't be called on the SnapshotReadCache.eventLoop, so there won't be a deadlock. + eventLoop.submit(this::clearOverloadedTask).join(); + } + eventLoop.submit(this::tryLoad); + return task.replayCf.whenCompleteAsync((nil, ex) -> tryLoad(), eventLoop); + } + + @EventLoopSafe + private void tryLoad() { + for (; ; ) { + if (loadingTasks.size() >= maxInflightLoadingCount) { + break; + } + WalReplayTask task = waitingLoadTasks.peek(); + if (task == null) { + break; + } + if (time.nanoseconds() - task.timestampNanos > TASK_WAITING_TIMEOUT_NANOS) { + clearOverloadedTask(); + return; + } + waitingLoadTasks.poll(); + loadingTasks.add(task); + task.run(); + task.loadCf.whenCompleteAsync((rst, ex) -> tryPutIntoCache(), eventLoop); + } + } + + /** + * Clears all waiting tasks when the replay system is overloaded. + * This is triggered when tasks wait longer than TASK_WAITING_TIMEOUT_NANOS or waitingLoadTasks is full. + * All dropped tasks have their futures completed with null, and affected + * nodes are notified to commit their WAL to free up resources. + */ + @EventLoopSafe + private void clearOverloadedTask() { + // The WalReplay is overloaded, so we need to drain all tasks promptly. + Set nodeIds = new HashSet<>(); + int dropCount = 0; + for (; ; ) { + WalReplayTask task = waitingLoadTasks.poll(); + if (task == null) { + break; + } + nodeIds.add(task.wal.metadata().nodeId()); + task.loadCf.complete(null); + task.replayCf.complete(null); + if (task.walRecords != null) { + task.walRecords.forEach(StreamRecordBatch::release); + } + dropCount++; + } + nodeIds.forEach(cacheFreeListener::notifyListener); + LOGGER.warn("wal replay is overloaded, drop all {} waiting tasks and request nodes={} to commit", dropCount, nodeIds); + } + + @EventLoopSafe + private void tryPutIntoCache() { + for (; ; ) { + WalReplayTask task = loadingTasks.peek(); + if (task == null || !task.loadCf.isDone()) { + break; + } + loadingTasks.poll(); + put(CloseableIterator.wrap(task.records.iterator())); + task.replayCf.complete(null); + } + } + + } + + class WalReplayTask { + final long timestampNanos = time.nanoseconds(); + final WriteAheadLog wal; + final RecordOffset startOffset; + final RecordOffset endOffset; + final List walRecords; + final CompletableFuture loadCf; + final CompletableFuture replayCf = new CompletableFuture<>(); + final List records = new ArrayList<>(); + + public WalReplayTask(WriteAheadLog wal, RecordOffset startOffset, RecordOffset endOffset, + List walRecords) { + this.wal = wal; + this.startOffset = startOffset; + this.endOffset = endOffset; + this.walRecords = walRecords; + this.loadCf = new CompletableFuture<>(); + loadCf.whenComplete((rst, ex) -> { + if (ex != null) { + LOGGER.error("Replay WAL [{}, {}) fail, wal={}", startOffset, endOffset, wal, ex); + } + }); + } + + public void run() { + long startNanos = time.nanoseconds(); + CompletableFuture> walRecordsCf = walRecords != null ? + CompletableFuture.completedFuture(walRecords) : wal.get(startOffset, endOffset); + walRecordsCf.thenCompose(walRecords -> { + long readWalDoneNanos = time.nanoseconds(); + READ_WAL_LATENCY.record(readWalDoneNanos - startNanos); + List> cfList = new ArrayList<>(walRecords.size()); + for (StreamRecordBatch walRecord : walRecords) { + if (walRecord.getCount() >= 0) { + cfList.add(CompletableFuture.completedFuture(walRecord)); + } else { + cfList.add(linkRecordDecoder.decode(walRecord)); + } + } + return CompletableFuture.allOf(cfList.toArray(new CompletableFuture[0])).whenComplete((rst, ex) -> { + DECODE_LATENCY.record(time.nanoseconds() - readWalDoneNanos); + if (ex != null) { + loadCf.completeExceptionally(ex); + // release other success record + cfList.forEach(cf -> cf.thenAccept(StreamRecordBatch::release)); + return; + } + records.addAll(cfList.stream().map(CompletableFuture::join).toList()); + ListIterator it = records.listIterator(); + while (it.hasNext()) { + StreamRecordBatch record = it.next(); + try { + // Copy the record to the SeqAlloc to reduce fragmentation. + it.set(StreamRecordBatch.parse(record.encoded(), true, ENCODE_ALLOC)); + } finally { + record.release(); + } + } + loadCf.complete(null); + }); + }).whenComplete((rst, ex) -> { + if (ex != null) { + loadCf.completeExceptionally(ex); + } + }); + } + } + + class ObjectReplay { + private final AtomicLong inflightLoadBytes = new AtomicLong(); + private final Queue waitingLoadingTasks = new LinkedList<>(); + private final Queue loadingTasks = new LinkedList<>(); + + public synchronized CompletableFuture replay(List objects) { + if (objects.isEmpty()) { + throw new IllegalArgumentException("The objects is an empty list"); + } + CompletableFuture cf = new CompletableFuture<>(); + eventLoop.execute(() -> { + ObjectReplayTask task = null; + for (S3ObjectMetadata object : objects) { + task = new ObjectReplayTask(ObjectReader.reader(object, objectStorage)); + waitingLoadingTasks.add(task); + } + if (task != null) { + FutureUtil.propagate(task.cf, cf); + } + tryLoad(); + }); + return cf; + } + + @EventLoopSafe + private void tryLoad() { + for (; ; ) { + if (inflightLoadBytes.get() >= MAX_INFLIGHT_LOAD_BYTES) { + break; + } + ObjectReplayTask task = waitingLoadingTasks.peek(); + if (task == null) { + break; + } + waitingLoadingTasks.poll(); + loadingTasks.add(task); + task.run(); + inflightLoadBytes.addAndGet(task.reader.metadata().objectSize()); + } + } + + @EventLoopSafe + private void tryPutIntoCache() { + for (; ; ) { + ObjectReplayTask task = loadingTasks.peek(); + if (task == null) { + break; + } + if (task.putIntoCache()) { + loadingTasks.poll(); + task.close(); + inflightLoadBytes.addAndGet(-task.reader.metadata().objectSize()); + tryLoad(); + } else { + break; + } + } + } + } + + class ObjectReplayTask { + final ObjectReader reader; + final CompletableFuture cf; + Queue> blocks; + + public ObjectReplayTask(ObjectReader reader) { + this.reader = reader; + this.cf = new CompletableFuture<>(); + } + + public void run() { + reader.basicObjectInfo().thenAcceptAsync(info -> { + List blockIndexList = info.indexBlock().indexes(); + LinkedList> blocks = new LinkedList<>(); + blockIndexList.forEach(blockIndex -> { + CompletableFuture blockCf = reader.read(blockIndex); + blocks.add(blockCf); + blockCf.whenCompleteAsync((group, t) -> { + if (t != null) { + LOGGER.error("Failed to load object blocks {}", reader.metadata(), t); + } + objectReplay.tryPutIntoCache(); + }, eventLoop); + }); + this.blocks = blocks; + }, eventLoop).exceptionally(ex -> { + LOGGER.error("Failed to load object {}", reader.metadata(), ex); + cf.complete(null); + return null; + }); + } + + /** + * Put blocks' records into cache + * + * @return true if all the data is put into cache, false otherwise + */ + @EventLoopSafe + public boolean putIntoCache() { + if (blocks == null) { + return false; + } + for (; ; ) { + CompletableFuture blockCf = blocks.peek(); + if (blockCf == null) { + cf.complete(null); + return true; + } + if (!blockCf.isDone()) { + return false; + } + if (blockCf.isCompletedExceptionally() || blockCf.isCancelled()) { + blocks.poll(); + continue; + } + try (ObjectReader.DataBlockGroup block = blockCf.join()) { + put(block.iterator()); + } + blocks.poll(); + } + } + + public void close() { + reader.close(); + } + } + + class LogCacheBlockFreeListener implements LogCache.FreeListener { + private final List listeners = new CopyOnWriteArrayList<>(); + + @Override + public void onFree(List bounds) { + List streamIdList = bounds.stream().map(LogCache.StreamRangeBound::streamId).collect(Collectors.toList()); + Map streamMap = bounds.stream().collect(Collectors.toMap(LogCache.StreamRangeBound::streamId, Function.identity())); + List streamMetadataList = streamManager.getStreams(streamIdList).join(); + Set requestCommitNodes = new HashSet<>(); + for (StreamMetadata streamMetadata : streamMetadataList) { + LogCache.StreamRangeBound bound = streamMap.get(streamMetadata.streamId()); + if (bound.endOffset() > streamMetadata.endOffset()) { + requestCommitNodes.add(streamMetadata.nodeId()); + } + } + requestCommitNodes.forEach(this::notifyListener); + } + + public void notifyListener(int nodeId) { + listeners.forEach(listener -> FutureUtil.suppress(() -> listener.onEvent(new RequestCommitEvent(nodeId)), LOGGER)); + } + + public void addListener(EventListener listener) { + this.listeners.add(listener); + } + } + + public interface EventListener { + void onEvent(Event event); + } + + public interface Event { + } + + public static class RequestCommitEvent implements Event { + private final int nodeId; + + public RequestCommitEvent(int nodeId) { + this.nodeId = nodeId; + } + + public int nodeId() { + return nodeId; + } + + @Override + public String toString() { + return "RequestCommitEvent{" + + "nodeId=" + nodeId + + '}'; + } + } + +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlock.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlock.java index 7f3525daf8..0fa50b9648 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlock.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlock.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -16,6 +24,7 @@ import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.utils.CloseableIterator; import com.automq.stream.utils.Time; +import com.automq.stream.utils.threads.EventLoopSafe; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlockCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlockCache.java index 84244603a5..c6d6c654ad 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlockCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DataBlockCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -18,9 +26,11 @@ import com.automq.stream.s3.metrics.S3StreamMetricsManager; import com.automq.stream.s3.metrics.stats.StorageOperationStats; import com.automq.stream.s3.network.ThrottleStrategy; +import com.automq.stream.utils.AsyncSemaphore; import com.automq.stream.utils.FutureUtil; import com.automq.stream.utils.Time; import com.automq.stream.utils.threads.EventLoop; +import com.automq.stream.utils.threads.EventLoopSafe; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DefaultObjectReaderFactory.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DefaultObjectReaderFactory.java index 4ffaf7643a..6334647673 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DefaultObjectReaderFactory.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/DefaultObjectReaderFactory.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -16,29 +24,37 @@ import com.automq.stream.s3.metadata.S3ObjectMetadata; import com.automq.stream.s3.operator.ObjectStorage; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; + public class DefaultObjectReaderFactory implements ObjectReaderFactory { private static final int MAX_OBJECT_READER_SIZE = 100 * 1024 * 1024; // 100MB; private final ObjectReaderLRUCache objectReaders; - private final ObjectStorage objectStorage; + private final Supplier objectStorage; public DefaultObjectReaderFactory(ObjectStorage objectStorage) { + this(() -> objectStorage); + } + + public DefaultObjectReaderFactory(Supplier objectStorage) { this.objectReaders = new ObjectReaderLRUCache("ObjectReader", MAX_OBJECT_READER_SIZE); this.objectStorage = objectStorage; } @Override public synchronized ObjectReader get(S3ObjectMetadata metadata) { - ObjectReader objectReader = objectReaders.get(metadata.objectId()); - if (objectReader == null) { - objectReader = ObjectReader.reader(metadata, objectStorage); - objectReaders.put(metadata.objectId(), objectReader); - } - return objectReader.retain(); + AtomicReference objectReaderRef = new AtomicReference<>(); + objectReaders.inLockRun(() -> { + ObjectReader objectReader = objectReaders.computeIfAbsent(metadata.objectId(), k -> ObjectReader.reader(metadata, objectStorage.get())); + objectReader.retain(); + objectReaderRef.set(objectReader); + }); + return objectReaderRef.get(); } @Override public ObjectStorage getObjectStorage() { - return objectStorage; + return objectStorage.get(); } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/EventLoopSafe.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/EventLoopSafe.java deleted file mode 100644 index 15fa2b991e..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/EventLoopSafe.java +++ /dev/null @@ -1,18 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.cache.blockcache; - -/** - * All the methods in this class should be called from the same stream's eventLoop. - */ -public @interface EventLoopSafe { -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ObjectReaderFactory.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ObjectReaderFactory.java index 38b69c836a..7b2cb0e8ad 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ObjectReaderFactory.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ObjectReaderFactory.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ReadStatusChangeListener.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ReadStatusChangeListener.java index 19c58c71a7..796727b80c 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ReadStatusChangeListener.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/ReadStatusChangeListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReader.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReader.java index 9dd5d092a4..f221e9f177 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -26,7 +34,10 @@ import com.automq.stream.s3.objects.ObjectManager; import com.automq.stream.utils.FutureUtil; import com.automq.stream.utils.LogSuppressor; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Time; import com.automq.stream.utils.threads.EventLoop; +import com.automq.stream.utils.threads.EventLoopSafe; import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; @@ -54,7 +65,7 @@ public static final int GET_OBJECT_STEP = 4; private static final Logger LOGGER = LoggerFactory.getLogger(StreamReader.class); static final int READAHEAD_SIZE_UNIT = 1024 * 1024 / 2; - private static final int MAX_READAHEAD_SIZE = 32 * 1024 * 1024; + private static final int MAX_READAHEAD_SIZE = Systems.getEnvInt("AUTOMQ_MAX_READAHEAD_SIZE", 32 * 1024 * 1024); private static final long READAHEAD_RESET_COLD_DOWN_MILLS = TimeUnit.MINUTES.toMillis(1); private static final long READAHEAD_AVAILABLE_BYTES_THRESHOLD = 32L * 1024 * 1024; private static final LogSuppressor READAHEAD_RESET_LOG_SUPPRESSOR = new LogSuppressor(LOGGER, 30000); @@ -72,19 +83,27 @@ private final ObjectManager objectManager; private final ObjectReaderFactory objectReaderFactory; private final DataBlockCache dataBlockCache; + private final Time time; long nextReadOffset; private CompletableFuture inflightLoadIndexCf; private volatile CompletableFuture afterReadTryReadaheadCf; - private long lastAccessTimestamp = System.currentTimeMillis(); + private long lastAccessTimestamp; private boolean reading = false; private boolean closed = false; public StreamReader(long streamId, long nextReadOffset, EventLoop eventLoop, ObjectManager objectManager, ObjectReaderFactory objectReaderFactory, DataBlockCache dataBlockCache) { + this(streamId, nextReadOffset, eventLoop, objectManager, objectReaderFactory, dataBlockCache, Time.SYSTEM); + } + + public StreamReader(long streamId, long nextReadOffset, EventLoop eventLoop, ObjectManager objectManager, + ObjectReaderFactory objectReaderFactory, DataBlockCache dataBlockCache, Time time) { this.streamId = streamId; this.nextReadOffset = nextReadOffset; this.readahead = new Readahead(); + this.time = time; + this.lastAccessTimestamp = time.milliseconds(); this.eventLoop = eventLoop; this.objectManager = objectManager; @@ -101,7 +120,7 @@ public CompletableFuture read(long startOffset, long endOffset, i } reading = true; try { - return read(startOffset, endOffset, maxBytes, 1).whenComplete((rst, ex) -> reading = false); + return read(startOffset, endOffset, maxBytes, 2).whenComplete((rst, ex) -> reading = false); } catch (Throwable e) { reading = false; return FutureUtil.failedFuture(e); @@ -109,7 +128,7 @@ public CompletableFuture read(long startOffset, long endOffset, i } CompletableFuture read(long startOffset, long endOffset, int maxBytes, int leftRetries) { - lastAccessTimestamp = System.currentTimeMillis(); + lastAccessTimestamp = time.milliseconds(); ReadContext readContext = new ReadContext(); read0(readContext, startOffset, endOffset, maxBytes); CompletableFuture retCf = new CompletableFuture<>(); @@ -609,7 +628,7 @@ class Readahead { private int cacheMissCount; public void tryReadahead(boolean cacheMiss) { - if (System.currentTimeMillis() - resetTimestamp < READAHEAD_RESET_COLD_DOWN_MILLS) { + if (time.milliseconds() - resetTimestamp < READAHEAD_RESET_COLD_DOWN_MILLS) { // skip readahead when readahead is in cold down return; } @@ -652,7 +671,7 @@ public void tryReadahead(boolean cacheMiss) { public void reset() { requireReset = true; - resetTimestamp = System.currentTimeMillis(); + resetTimestamp = time.milliseconds(); } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReaders.java b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReaders.java index 9e7d3d64fd..70bd04ae5d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReaders.java +++ b/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/StreamReaders.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -18,7 +26,10 @@ import com.automq.stream.s3.trace.context.TraceContext; import com.automq.stream.utils.FutureUtil; import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.Time; import com.automq.stream.utils.threads.EventLoop; +import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -35,6 +46,7 @@ public class StreamReaders implements S3BlockCache { private static final long STREAM_READER_EXPIRED_MILLS = TimeUnit.MINUTES.toMillis(1); private static final long STREAM_READER_EXPIRED_CHECK_INTERVAL_MILLS = TimeUnit.MINUTES.toMillis(1); private final Cache[] caches; + private final Time time; private final DataBlockCache dataBlockCache; private final ObjectReaderFactory objectReaderFactory; @@ -43,11 +55,18 @@ public class StreamReaders implements S3BlockCache { public StreamReaders(long size, ObjectManager objectManager, ObjectStorage objectStorage, ObjectReaderFactory objectReaderFactory) { - this(size, objectManager, objectStorage, objectReaderFactory, Systems.CPU_CORES); + this(size, objectManager, objectStorage, objectReaderFactory, Systems.CPU_CORES, Time.SYSTEM); } public StreamReaders(long size, ObjectManager objectManager, ObjectStorage objectStorage, ObjectReaderFactory objectReaderFactory, int concurrency) { + this(size, objectManager, objectStorage, objectReaderFactory, concurrency, Time.SYSTEM); + } + + @SuppressWarnings("this-escape") + public StreamReaders(long size, ObjectManager objectManager, ObjectStorage objectStorage, + ObjectReaderFactory objectReaderFactory, int concurrency, Time time) { + this.time = time; EventLoop[] eventLoops = new EventLoop[concurrency]; for (int i = 0; i < concurrency; i++) { eventLoops[i] = new EventLoop("stream-reader-" + i); @@ -61,6 +80,11 @@ public StreamReaders(long size, ObjectManager objectManager, ObjectStorage objec this.objectReaderFactory = objectReaderFactory; this.objectManager = objectManager; this.objectStorage = objectStorage; + + Threads.COMMON_SCHEDULER.scheduleAtFixedRate(this::triggerExpiredStreamReaderCleanup, + STREAM_READER_EXPIRED_CHECK_INTERVAL_MILLS, + STREAM_READER_EXPIRED_CHECK_INTERVAL_MILLS, + TimeUnit.MILLISECONDS); } @Override @@ -70,6 +94,26 @@ public CompletableFuture read(TraceContext context, long streamId return cache.read(streamId, startOffset, endOffset, maxBytes); } + /** + * Get the total number of active StreamReaders across all caches. + * This method is intended for testing purposes only. + */ + @VisibleForTesting + int getActiveStreamReaderCount() { + int total = 0; + for (Cache cache : caches) { + total += cache.getStreamReaderCount(); + } + return total; + } + + @VisibleForTesting + void triggerExpiredStreamReaderCleanup() { + for (Cache cache : caches) { + cache.submitCleanupExpiredStreamReader(); + } + } + static class StreamReaderKey { final long streamId; final long startOffset; @@ -106,10 +150,11 @@ public String toString() { class Cache { private final EventLoop eventLoop; private final Map streamReaders = new HashMap<>(); - private long lastStreamReaderExpiredCheckTime = System.currentTimeMillis(); + private long lastStreamReaderExpiredCheckTime; public Cache(EventLoop eventLoop) { this.eventLoop = eventLoop; + this.lastStreamReaderExpiredCheckTime = time.milliseconds(); } public CompletableFuture read(long streamId, long startOffset, @@ -121,7 +166,7 @@ public CompletableFuture read(long streamId, long startOffset, StreamReaderKey key = new StreamReaderKey(streamId, startOffset); StreamReader streamReader = streamReaders.remove(key); if (streamReader == null) { - streamReader = new StreamReader(streamId, startOffset, eventLoop, objectManager, objectReaderFactory, dataBlockCache); + streamReader = new StreamReader(streamId, startOffset, eventLoop, objectManager, objectReaderFactory, dataBlockCache, time); } StreamReader finalStreamReader = streamReader; CompletableFuture streamReadCf = streamReader.read(startOffset, endOffset, maxBytes) @@ -142,8 +187,21 @@ public CompletableFuture read(long streamId, long startOffset, return cf; } + private void submitCleanupExpiredStreamReader() { + eventLoop.execute(this::cleanupExpiredStreamReader); + } + + /** + * Get the number of StreamReaders in this cache. + * This method is intended for testing purposes only. + */ + @VisibleForTesting + int getStreamReaderCount() { + return streamReaders.size(); + } + private void cleanupExpiredStreamReader() { - long now = System.currentTimeMillis(); + long now = time.milliseconds(); if (now > lastStreamReaderExpiredCheckTime + STREAM_READER_EXPIRED_CHECK_INTERVAL_MILLS) { lastStreamReaderExpiredCheckTime = now; Iterator> it = streamReaders.entrySet().iterator(); diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactOperations.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactOperations.java index 0711dfbbfe..bb2c033119 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactOperations.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactOperations.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactResult.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactResult.java index 5772d936ff..1200644eaf 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactResult.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactResult.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionAnalyzer.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionAnalyzer.java index 347c4f355e..91e8ca2b3e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionAnalyzer.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionAnalyzer.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionConstants.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionConstants.java index efaa38e849..42cd87e2f9 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionConstants.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionConstants.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionManager.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionManager.java index 19851eccef..a6bbd6d48c 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionManager.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; @@ -30,8 +38,10 @@ import com.automq.stream.s3.objects.ObjectManager; import com.automq.stream.s3.objects.ObjectStreamRange; import com.automq.stream.s3.objects.StreamObject; +import com.automq.stream.s3.operator.LocalFileObjectStorage; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.utils.FutureUtil; import com.automq.stream.utils.LogContext; import com.automq.stream.utils.ThreadUtils; import com.automq.stream.utils.Threads; @@ -45,7 +55,6 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Collections; -import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -56,19 +65,17 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import io.github.bucket4j.Bucket; -import io.netty.util.concurrent.DefaultThreadFactory; import static com.automq.stream.s3.metadata.ObjectUtils.NOOP_OBJECT_ID; public class CompactionManager { - private static final int MIN_COMPACTION_DELAY_MS = 60000; + private static final int MIN_COMPACTION_DELAY_MS = 10000; // Max refill rate for Bucket: 1 token per nanosecond private static final int MAX_THROTTLE_BYTES_PER_SEC = 1000000000; private final Logger logger; @@ -125,8 +132,8 @@ public CompactionManager(Config config, ObjectManager objectManager, StreamManag ThreadUtils.createThreadFactory("s3-data-block-reader-bucket-cb-%d", true), logger, true, false); this.utilityScheduledExecutor = Threads.newSingleThreadScheduledExecutor( ThreadUtils.createThreadFactory("compaction-utility-executor-%d", true), logger, true, false); - this.compactThreadPool = Executors.newFixedThreadPool(1, new DefaultThreadFactory("object-compaction-manager")); - this.forceSplitThreadPool = Executors.newFixedThreadPool(1, new DefaultThreadFactory("force-split-executor")); + this.compactThreadPool = Threads.newFixedThreadPoolWithMonitor(1, "object-compaction-manager", true, logger); + this.forceSplitThreadPool = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(1, "force-split-executor", true, logger); this.running.set(true); S3StreamMetricsManager.registerCompactionDelayTimeSuppler(() -> compactionDelayTime); this.logger.info("Compaction manager initialized with config: compactionInterval: {} min, compactionCacheSize: {} bytes, " + @@ -146,9 +153,19 @@ public void start() { this.compactionDelayTime = 0; return; } - data.sort(Comparator.comparingLong(S3ObjectMetadata::committedTimestamp)); - this.compactionDelayTime = System.currentTimeMillis() - data.get(0).committedTimestamp(); - }).join(), 1, 1, TimeUnit.MINUTES); + long minCommittedTimestamp = Long.MAX_VALUE; + for (S3ObjectMetadata metadata : data) { + long ts = metadata.committedTimestamp(); + if (ts < minCommittedTimestamp) { + minCommittedTimestamp = ts; + } + } + if (minCommittedTimestamp == Long.MAX_VALUE) { + this.compactionDelayTime = 0; + return; + } + this.compactionDelayTime = System.currentTimeMillis() - minCommittedTimestamp; + }).join(), (long) this.compactionInterval * 2, 1, TimeUnit.MINUTES); } void scheduleNextCompaction(long delayMillis) { @@ -205,16 +222,7 @@ public void shutdown() { } private void shutdownAndAwaitTermination(ExecutorService executor, int timeout, TimeUnit timeUnit) { - executor.shutdown(); - try { - if (!executor.awaitTermination(timeout, timeUnit)) { - executor.shutdownNow(); - } - } catch (InterruptedException ex) { - executor.shutdownNow(); - // Preserve interrupt status - Thread.currentThread().interrupt(); - } + ThreadUtils.shutdownExecutor(executor, timeout, timeUnit, logger); } public CompletableFuture compact() { @@ -223,6 +231,10 @@ public CompletableFuture compact() { if (objectMetadataList.isEmpty()) { return CompletableFuture.completedFuture(null); } + if (objectMetadataList.stream().anyMatch(o -> o.bucket() == LocalFileObjectStorage.BUCKET_ID)) { + logger.info("Skip the compaction, because main storage circuit breaker isn't in closed status"); + return CompletableFuture.completedFuture(null); + } updateStreamDataBlockMap(objectMetadataList); List streamIds; try { @@ -276,22 +288,16 @@ private void compact(List streamMetadataList, long totalSize = objectsToForceSplit.stream().mapToLong(S3ObjectMetadata::objectSize).sum(); totalSize += objectsToCompact.stream().mapToLong(S3ObjectMetadata::objectSize).sum(); // throttle compaction read to half of compaction interval because of write overhead - int expectCompleteTime = compactionInterval - 1 /* ahead 1min*/; - long expectReadBytesPerSec; - if (expectCompleteTime > 0) { - expectReadBytesPerSec = totalSize / expectCompleteTime / 60; - if (expectReadBytesPerSec < MAX_THROTTLE_BYTES_PER_SEC) { - compactionBucket = Bucket.builder().addLimit(limit -> limit - .capacity(expectReadBytesPerSec) - .refillIntervally(expectReadBytesPerSec, Duration.ofSeconds(1))).build(); - logger.info("Throttle compaction read to {} bytes/s, expect to complete in no less than {}min", - expectReadBytesPerSec, expectCompleteTime); - } else { - logger.warn("Compaction throttle rate {} bytes/s exceeds bucket refill limit, there will be no throttle for compaction this time", expectReadBytesPerSec); - compactionBucket = null; - } + int expectCompleteTime = Math.max(compactionInterval - 1, 1) /* ahead 1min*/; + long expectReadBytesPerSec = Math.max(expectCompleteTime * 60L, totalSize / expectCompleteTime / 60); + if (expectReadBytesPerSec < MAX_THROTTLE_BYTES_PER_SEC) { + compactionBucket = Bucket.builder().addLimit(limit -> limit + .capacity(expectReadBytesPerSec) + .refillIntervally(expectReadBytesPerSec, Duration.ofSeconds(1))).build(); + logger.info("Throttle compaction read to {} bytes/s, expect to complete in no less than {}min", + expectReadBytesPerSec, expectCompleteTime); } else { - logger.warn("Compaction interval {}min is too small, there will be no throttle for compaction this time", compactionInterval); + logger.warn("Compaction throttle rate {} bytes/s exceeds bucket refill limit, there will be no throttle for compaction this time", expectReadBytesPerSec); compactionBucket = null; } @@ -697,6 +703,7 @@ Map> convertS3Objects(List str >= TimeUnit.MINUTES.toMillis(this.forceSplitObjectPeriod)))); } + @SuppressWarnings("checkstyle:CyclomaticComplexity") void executeCompactionPlans(CommitStreamSetObjectRequest request, List compactionPlans, List s3ObjectMetadata) throws CompletionException { @@ -743,15 +750,35 @@ void executeCompactionPlans(CommitStreamSetObjectRequest request, List uploader.forceUploadStreamSetObject()) - .exceptionally(ex -> { - uploader.release().thenAccept(v -> { - for (CompactedObject compactedObject : compactionPlan.compactedObjects()) { - compactedObject.streamDataBlocks().forEach(StreamDataBlock::release); - } - }).join(); - throw new IllegalStateException("Error while uploading compaction objects", ex); + compactionCf = new CompletableFuture<>(); + CompletableFuture.allOf(cfList.toArray(new CompletableFuture[0])) + .whenComplete((v, uploadException) -> { + if (uploadException != null) { + logger.error("Error while uploading compaction objects", uploadException); + } + FutureUtil.exec(uploader::forceUploadStreamSetObject, logger, "force upload sso") + .whenComplete((vv, forceUploadException) -> { + if (forceUploadException != null) { + logger.error("Error while force uploading stream set object", uploadException); + } + if (uploadException != null || forceUploadException != null) { + FutureUtil.exec(uploader::release, logger, "release uploader").whenComplete((vvv, releaseException) -> { + if (releaseException != null) { + logger.error("Unexpected exception while release uploader"); + } + for (CompactedObject compactedObject : compactionPlan.compactedObjects()) { + compactedObject.streamDataBlocks().forEach(StreamDataBlock::release); + } + if (uploadException != null) { + compactionCf.completeExceptionally(new CompletionException("Uploading failed", uploadException)); + } else { + compactionCf.completeExceptionally(new CompletionException("Force uploading sso failed", forceUploadException)); + } + }); + } else { + compactionCf.complete(null); + } + }); }); } try { diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionPlan.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionPlan.java index 469e7c11bf..5b56540f8b 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionPlan.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionPlan.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionStats.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionStats.java index 9f99899c97..8721f58d8e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionUploader.java b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionUploader.java index 4ec62d31b6..9687eae397 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionUploader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/CompactionUploader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/StreamObjectCompactor.java b/s3stream/src/main/java/com/automq/stream/s3/compact/StreamObjectCompactor.java index ff8ce6dcd7..0e4f610bbb 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/StreamObjectCompactor.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/StreamObjectCompactor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; @@ -30,6 +38,7 @@ import com.automq.stream.s3.objects.ObjectAttributes; import com.automq.stream.s3.objects.ObjectManager; import com.automq.stream.s3.objects.ObjectStreamRange; +import com.automq.stream.s3.operator.LocalFileObjectStorage; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorage.ObjectPath; import com.automq.stream.s3.operator.ObjectStorage.WriteOptions; @@ -224,6 +233,9 @@ static boolean checkObjectGroupCouldBeCompact(List objectGroup if (objectGroup.size() == 1 && SKIP_COMPACTION_TYPE_WHEN_ONE_OBJECT_IN_GROUP.contains(compactionType)) { return false; } + if (objectGroup.stream().anyMatch(o -> o.bucket() == LocalFileObjectStorage.BUCKET_ID)) { + return false; + } if (CLEANUP_V1.equals(compactionType)) { S3ObjectMetadata metadata = objectGroup.get(0); if (ObjectAttributes.from(metadata.attributes()).type() != Composite) { diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObject.java b/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObject.java index 3f4fbaa24a..403fdeeb66 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObject.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObject.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObjectBuilder.java b/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObjectBuilder.java index 2babc5c584..e1ce02ab28 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObjectBuilder.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactedObjectBuilder.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactionType.java b/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactionType.java index d9b46e1c16..6e95c452e9 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactionType.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/objects/CompactionType.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockReader.java b/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockReader.java index ce766cca9b..3760a210f8 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockReader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockReader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockWriter.java b/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockWriter.java index 45c3d6f256..fc80b9ec04 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockWriter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/operator/DataBlockWriter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.operator; @@ -85,10 +93,14 @@ public CompletableFuture forceUpload() { private void uploadWaitingList() { CompositeByteBuf buf = groupWaitingBlocks(); List blocks = new LinkedList<>(waitingUploadBlocks); - writer.write(buf).thenAccept(v -> { + writer.write(buf).whenComplete((v, ex) -> { for (StreamDataBlock block : blocks) { waitingUploadBlockCfs.computeIfPresent(block, (k, cf) -> { - cf.complete(null); + if (ex != null) { + cf.completeExceptionally(ex); + } else { + cf.complete(null); + } return null; }); } @@ -128,8 +140,7 @@ public CompletableFuture close() { private CompositeByteBuf groupWaitingBlocks() { CompositeByteBuf buf = ByteBufAlloc.compositeByteBuffer(); for (StreamDataBlock block : waitingUploadBlocks) { - buf.addComponent(true, block.getDataCf().join()); - block.releaseRef(); + buf.addComponent(true, block.getAndReleaseData()); completedBlocks.add(block); nextDataBlockPosition += block.getBlockSize(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/utils/CompactionUtils.java b/s3stream/src/main/java/com/automq/stream/s3/compact/utils/CompactionUtils.java index 86e9c001bf..6214320679 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/utils/CompactionUtils.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/utils/CompactionUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.utils; @@ -36,29 +44,6 @@ import java.util.stream.Collectors; public class CompactionUtils { - public static List buildObjectStreamRange(List streamDataBlocks) { - List objectStreamRanges = new ArrayList<>(); - ObjectStreamRange currObjectStreamRange = null; - for (StreamDataBlock streamDataBlock : streamDataBlocks) { - if (currObjectStreamRange == null) { - currObjectStreamRange = new ObjectStreamRange(streamDataBlock.getStreamId(), -1L, - streamDataBlock.getStartOffset(), streamDataBlock.getEndOffset(), streamDataBlock.getBlockSize()); - } else { - if (currObjectStreamRange.getStreamId() == streamDataBlock.getStreamId()) { - currObjectStreamRange.setEndOffset(streamDataBlock.getEndOffset()); - currObjectStreamRange.setSize(currObjectStreamRange.getSize() + streamDataBlock.getBlockSize()); - } else { - objectStreamRanges.add(currObjectStreamRange); - currObjectStreamRange = new ObjectStreamRange(streamDataBlock.getStreamId(), -1L, - streamDataBlock.getStartOffset(), streamDataBlock.getEndOffset(), streamDataBlock.getBlockSize()); - } - } - } - if (currObjectStreamRange != null) { - objectStreamRanges.add(currObjectStreamRange); - } - return objectStreamRanges; - } // test only public static Map> blockWaitObjectIndices(List streamMetadataList, List objectMetadataList, diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByLimitPredicate.java b/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByLimitPredicate.java index ed2a723a43..cf098a5745 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByLimitPredicate.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByLimitPredicate.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.utils; diff --git a/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByOffsetPredicate.java b/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByOffsetPredicate.java index cc5b8d383e..bf0f2e7c5c 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByOffsetPredicate.java +++ b/s3stream/src/main/java/com/automq/stream/s3/compact/utils/GroupByOffsetPredicate.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact.utils; diff --git a/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java b/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java index f16e985876..1114c21da2 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java +++ b/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java @@ -1,23 +1,33 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.context; import com.automq.stream.s3.trace.context.TraceContext; +import io.netty.buffer.ByteBuf; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; public class AppendContext extends TraceContext { public static final AppendContext DEFAULT = new AppendContext(); + private ByteBuf linkRecord; public AppendContext() { super(false, null, null); @@ -30,4 +40,12 @@ public AppendContext(TraceContext context) { public AppendContext(boolean isTraceEnabled, Tracer tracer, Context currentContext) { super(isTraceEnabled, tracer, currentContext); } + + public void linkRecord(ByteBuf record) { + this.linkRecord = record; + } + + public ByteBuf linkRecord() { + return this.linkRecord; + } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java b/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java index be0256b221..5a4dd3aa4d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java +++ b/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.context; diff --git a/s3stream/src/main/java/com/automq/stream/s3/exceptions/AutoMQException.java b/s3stream/src/main/java/com/automq/stream/s3/exceptions/AutoMQException.java index cf14ce1a79..8b6fafb7d9 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/exceptions/AutoMQException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/exceptions/AutoMQException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.exceptions; diff --git a/s3stream/src/main/java/com/automq/stream/s3/exceptions/BlockNotContinuousException.java b/s3stream/src/main/java/com/automq/stream/s3/exceptions/BlockNotContinuousException.java index a232de229e..7215471444 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/exceptions/BlockNotContinuousException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/exceptions/BlockNotContinuousException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.exceptions; diff --git a/s3stream/src/main/java/com/automq/stream/s3/exceptions/CompactedObjectsNotFoundException.java b/s3stream/src/main/java/com/automq/stream/s3/exceptions/CompactedObjectsNotFoundException.java new file mode 100644 index 0000000000..976307292c --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/exceptions/CompactedObjectsNotFoundException.java @@ -0,0 +1,23 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.exceptions; + +public class CompactedObjectsNotFoundException extends AutoMQException { +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/exceptions/IndexBlockParseException.java b/s3stream/src/main/java/com/automq/stream/s3/exceptions/IndexBlockParseException.java index 7238a03559..24768eda69 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/exceptions/IndexBlockParseException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/exceptions/IndexBlockParseException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.exceptions; diff --git a/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotCommittedException.java b/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotCommittedException.java new file mode 100644 index 0000000000..4fbf5aef9b --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotCommittedException.java @@ -0,0 +1,23 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.exceptions; + +public class ObjectNotCommittedException extends AutoMQException { +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotExistException.java b/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotExistException.java index 9380a208e8..2212b0b8b7 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotExistException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/exceptions/ObjectNotExistException.java @@ -1,24 +1,40 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.exceptions; public class ObjectNotExistException extends AutoMQException { + public ObjectNotExistException() { + } + public ObjectNotExistException(long objectId) { super("Object not exist: " + objectId); } + public ObjectNotExistException(String msg) { + super(msg); + } + + @SuppressWarnings("this-escape") public ObjectNotExistException(Throwable cause) { - super(cause); + super(cause.getMessage(), cause); } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/Failover.java b/s3stream/src/main/java/com/automq/stream/s3/failover/Failover.java index 33d6cb24cc..c2bc315874 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/Failover.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/Failover.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; @@ -27,9 +35,6 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; -import static com.automq.stream.s3.Constants.NOOP_EPOCH; -import static com.automq.stream.s3.Constants.NOOP_NODE_ID; - /** * To perform a Delta WAL failover, follow these steps: * 1. Ensure the old node stops writing to the delta WAL. @@ -63,8 +68,6 @@ public CompletableFuture failover(FailoverRequest request) { class FailoverTask { private final FailoverRequest request; - private int nodeId = NOOP_NODE_ID; - private long epoch = NOOP_EPOCH; public FailoverTask(FailoverRequest request) { this.request = request; @@ -72,8 +75,13 @@ public FailoverTask(FailoverRequest request) { public FailoverResponse failover() throws Throwable { LOGGER.info("failover start {}", request); + int nodeId = request.getNodeId(); + long nodeEpoch = request.getNodeEpoch(); + FailoverResponse resp = new FailoverResponse(); - resp.setNodeId(request.getNodeId()); + resp.setNodeId(nodeId); + resp.setEpoch(nodeEpoch); + // fence the device to ensure the old node stops writing to the delta WAL // recover WAL data and upload to S3 WriteAheadLog wal = factory.getWal(request); @@ -83,18 +91,19 @@ public FailoverResponse failover() throws Throwable { LOGGER.info("fail over empty wal {}", request); return resp; } + try { WALMetadata metadata = wal.metadata(); - this.nodeId = metadata.nodeId(); - this.epoch = metadata.epoch(); - if (nodeId != request.getNodeId()) { + if (nodeId != metadata.nodeId()) { throw new IllegalArgumentException(String.format("nodeId mismatch, request=%s, wal=%s", request, metadata)); } - resp.setNodeId(nodeId); - resp.setEpoch(epoch); - Logger taskLogger = new LogContext(String.format("[Failover nodeId=%s epoch=%s]", nodeId, epoch)).logger(FailoverTask.class); - StreamManager streamManager = factory.getStreamManager(nodeId, epoch); - ObjectManager objectManager = factory.getObjectManager(nodeId, epoch); + if (nodeEpoch < metadata.epoch()) { + throw new IllegalStateException(String.format("epoch mismatch, request=%s, wal=%s", request, metadata)); + } + + Logger taskLogger = new LogContext(String.format("[Failover nodeId=%s epoch=%s]", nodeId, nodeEpoch)).logger(FailoverTask.class); + StreamManager streamManager = factory.getStreamManager(nodeId, nodeEpoch); + ObjectManager objectManager = factory.getObjectManager(nodeId, nodeEpoch); LOGGER.info("failover recover {}", request); walRecover.recover(wal, streamManager, objectManager, taskLogger); } finally { @@ -108,8 +117,6 @@ public FailoverResponse failover() throws Throwable { public String toString() { return "FailoverTask{" + "request=" + request + - ", nodeId=" + nodeId + - ", epoch=" + epoch + '}'; } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverFactory.java b/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverFactory.java index 9e93313dae..9e851515db 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverFactory.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverFactory.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverRequest.java b/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverRequest.java index ce459b9727..b5fc06cce1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverRequest.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverRequest.java @@ -1,20 +1,30 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; public class FailoverRequest { private int nodeId; + private long nodeEpoch; private String volumeId; private String device; + private String kraftWalConfigs; public int getNodeId() { return nodeId; @@ -24,6 +34,14 @@ public void setNodeId(int nodeId) { this.nodeId = nodeId; } + public long getNodeEpoch() { + return nodeEpoch; + } + + public void setNodeEpoch(long nodeEpoch) { + this.nodeEpoch = nodeEpoch; + } + public String getVolumeId() { return volumeId; } @@ -40,12 +58,22 @@ public void setDevice(String device) { this.device = device; } + public String getKraftWalConfigs() { + return kraftWalConfigs; + } + + public void setKraftWalConfigs(String kraftWalConfigs) { + this.kraftWalConfigs = kraftWalConfigs; + } + @Override public String toString() { return "FailoverRequest{" + "nodeId=" + nodeId + + ", nodeEpoch=" + nodeEpoch + ", volumeId='" + volumeId + '\'' + ", device='" + device + '\'' + + ", kraftWalConfigs='" + kraftWalConfigs + '\'' + '}'; } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverResponse.java b/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverResponse.java index 541dec4647..6096bec208 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverResponse.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/FailoverResponse.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/ForceCloseStorageFailureHandler.java b/s3stream/src/main/java/com/automq/stream/s3/failover/ForceCloseStorageFailureHandler.java index a4b9d9cb0e..f632c9b504 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/ForceCloseStorageFailureHandler.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/ForceCloseStorageFailureHandler.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/HaltStorageFailureHandler.java b/s3stream/src/main/java/com/automq/stream/s3/failover/HaltStorageFailureHandler.java index 39c3553818..f3ebd0df4b 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/HaltStorageFailureHandler.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/HaltStorageFailureHandler.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandler.java b/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandler.java index 07f8451fee..aadb0032cb 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandler.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandler.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandlerChain.java b/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandlerChain.java index 2c6d09135b..882145062f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandlerChain.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/StorageFailureHandlerChain.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/failover/WALRecover.java b/s3stream/src/main/java/com/automq/stream/s3/failover/WALRecover.java index d0b361571c..1b3b3034c2 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/failover/WALRecover.java +++ b/s3stream/src/main/java/com/automq/stream/s3/failover/WALRecover.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.failover; diff --git a/s3stream/src/main/java/com/automq/stream/s3/index/LocalStreamRangeIndexCache.java b/s3stream/src/main/java/com/automq/stream/s3/index/LocalStreamRangeIndexCache.java index 52f1a7f654..89961be160 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/index/LocalStreamRangeIndexCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/index/LocalStreamRangeIndexCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; @@ -22,7 +30,7 @@ import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorage.ReadOptions; import com.automq.stream.utils.Systems; -import com.automq.stream.utils.ThreadUtils; +import com.automq.stream.utils.Threads; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -40,7 +48,6 @@ import java.util.Set; import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -62,8 +69,8 @@ public class LocalStreamRangeIndexCache implements S3StreamClient.StreamLifeCycl private final ReadWriteLock lock = new ReentrantReadWriteLock(); private final Lock readLock = lock.readLock(); private final Lock writeLock = lock.writeLock(); - private final ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor( - ThreadUtils.createThreadFactory("upload-index", true)); + private final ScheduledExecutorService executorService = + Threads.newSingleThreadScheduledExecutor("upload-index", true, LOGGER); private final Queue> uploadQueue = new LinkedList<>(); private final CompletableFuture initCf = new CompletableFuture<>(); private final AtomicBoolean pruned = new AtomicBoolean(false); @@ -73,7 +80,17 @@ public class LocalStreamRangeIndexCache implements S3StreamClient.StreamLifeCycl private CompletableFuture uploadCf = CompletableFuture.completedFuture(null); private long lastUploadTime = 0L; - public LocalStreamRangeIndexCache() { + private LocalStreamRangeIndexCache() { + + } + + public static LocalStreamRangeIndexCache create() { + LocalStreamRangeIndexCache cache = new LocalStreamRangeIndexCache(); + cache.completeInitialization(); + return cache; + } + + private void completeInitialization() { S3StreamMetricsManager.registerLocalStreamRangeIndexCacheSizeSupplier(this::totalSize); S3StreamMetricsManager.registerLocalStreamRangeIndexCacheStreamNumSupplier(() -> { readLock.lock(); @@ -108,7 +125,7 @@ public synchronized CompletableFuture uploadOnStreamClose() { lastUploadTime = now; LOGGER.info("Upload local index cache on stream close"); } - return uploadCf; + return uploadCf.orTimeout(1, TimeUnit.SECONDS); } CompletableFuture initCf() { @@ -317,36 +334,24 @@ public CompletableFuture compact(Map> rangeInde return exec(() -> { writeLock.lock(); try { - if (rangeIndexMap == null || rangeIndexMap.isEmpty()) { - Iterator> iterator = streamRangeIndexMap.entrySet().iterator(); - while (iterator.hasNext()) { - Map.Entry entry = iterator.next(); - totalSize += entry.getValue().compact(null, compactedObjectIds); - if (entry.getValue().length() == 0) { - iterator.remove(); - } - } - return null; - } - for (Map.Entry> entry : rangeIndexMap.entrySet()) { + // compact existing stream range index + Iterator> iterator = streamRangeIndexMap.entrySet().iterator(); + while (iterator.hasNext()) { + Map.Entry entry = iterator.next(); long streamId = entry.getKey(); - Optional rangeIndex = entry.getValue(); - streamRangeIndexMap.compute(streamId, (k, v) -> { - if (v == null) { - v = new SparseRangeIndex(COMPACT_NUM); - } - totalSize += v.compact(rangeIndex.orElse(null), compactedObjectIds); - if (v.length() == 0) { - // remove stream with empty index - return null; - } - return v; - }); + RangeIndex newRangeIndex = null; + if (rangeIndexMap.containsKey(streamId)) { + newRangeIndex = rangeIndexMap.get(streamId).orElse(null); + } + totalSize += entry.getValue().compact(newRangeIndex, compactedObjectIds); + if (entry.getValue().length() == 0) { + iterator.remove(); + } } + return null; } finally { writeLock.unlock(); } - return null; }); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/index/NodeRangeIndexCache.java b/s3stream/src/main/java/com/automq/stream/s3/index/NodeRangeIndexCache.java index 78f2c70a02..27cdac7498 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/index/NodeRangeIndexCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/index/NodeRangeIndexCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; diff --git a/s3stream/src/main/java/com/automq/stream/s3/index/RangeIndex.java b/s3stream/src/main/java/com/automq/stream/s3/index/RangeIndex.java index a6548a2e68..c151ca4c25 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/index/RangeIndex.java +++ b/s3stream/src/main/java/com/automq/stream/s3/index/RangeIndex.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; diff --git a/s3stream/src/main/java/com/automq/stream/s3/index/SparseRangeIndex.java b/s3stream/src/main/java/com/automq/stream/s3/index/SparseRangeIndex.java index 2260b3497c..87be91257f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/index/SparseRangeIndex.java +++ b/s3stream/src/main/java/com/automq/stream/s3/index/SparseRangeIndex.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; diff --git a/s3stream/src/main/java/com/automq/stream/s3/memory/MemoryMetadataManager.java b/s3stream/src/main/java/com/automq/stream/s3/memory/MemoryMetadataManager.java index 72e12177a9..b34d8e6895 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/memory/MemoryMetadataManager.java +++ b/s3stream/src/main/java/com/automq/stream/s3/memory/MemoryMetadataManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.memory; @@ -26,6 +34,7 @@ import com.automq.stream.s3.objects.StreamObject; import com.automq.stream.s3.streams.StreamCloseHook; import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.s3.streams.StreamMetadataListener; import org.apache.commons.lang3.tuple.Pair; @@ -250,6 +259,11 @@ public CompletableFuture> getStreams(List streamIds) return CompletableFuture.completedFuture(streamIds.stream().map(streams::get).filter(Objects::nonNull).collect(Collectors.toList())); } + @Override + public StreamMetadataListener.Handle addMetadataListener(long streamId, StreamMetadataListener listener) { + throw new UnsupportedOperationException(); + } + @Override public synchronized CompletableFuture createStream(Map tags) { long streamId = streamIdAlloc.getAndIncrement(); diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/ObjectUtils.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/ObjectUtils.java index b3a13b5401..3b5a29f792 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/ObjectUtils.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/ObjectUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectMetadata.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectMetadata.java index 2c0010b22c..92894b1a95 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectMetadata.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectMetadata.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectType.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectType.java index f6afbeccaa..aac5594e90 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectType.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/S3ObjectType.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/S3StreamConstant.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/S3StreamConstant.java index be36e9a346..695b6dbb0f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/S3StreamConstant.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/S3StreamConstant.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamMetadata.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamMetadata.java index 86e7b325ac..b46013b895 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamMetadata.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamMetadata.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; @@ -17,6 +25,7 @@ public class StreamMetadata { private long startOffset; private long endOffset; private StreamState state; + private int nodeId = -1; @SuppressWarnings("unused") public StreamMetadata() { @@ -70,6 +79,14 @@ public void state(StreamState state) { this.state = state; } + public int nodeId() { + return nodeId; + } + + public void nodeId(int nodeId) { + this.nodeId = nodeId; + } + @Override public String toString() { return "StreamMetadata{" + @@ -78,6 +95,7 @@ public String toString() { ", startOffset=" + startOffset + ", endOffset=" + endOffset + ", state=" + state + + ", nodeId=" + nodeId + '}'; } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamOffsetRange.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamOffsetRange.java index bc5a13f156..9fee0d104d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamOffsetRange.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamOffsetRange.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamState.java b/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamState.java index 43d91ca07e..6edf5b81f6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamState.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metadata/StreamState.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metadata; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/AttributesUtils.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/AttributesUtils.java index d12ae6bbd7..4c97e9c488 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/AttributesUtils.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/AttributesUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/Metrics.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/Metrics.java new file mode 100644 index 0000000000..5bc7d6f9ce --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/Metrics.java @@ -0,0 +1,406 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.metrics; + +import com.automq.stream.s3.metrics.wrapper.DeltaHistogram; + +import java.util.List; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.LongCounter; +import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.metrics.ObservableDoubleGauge; +import io.opentelemetry.api.metrics.ObservableDoubleMeasurement; +import io.opentelemetry.api.metrics.ObservableLongGauge; +import io.opentelemetry.api.metrics.ObservableLongMeasurement; +import io.opentelemetry.context.Context; + +public class Metrics { + private static final Metrics INSTANCE = new Metrics(); + private Meter meter; + private MetricsConfig globalConfig; + private final Queue waitingSetups = new ConcurrentLinkedQueue<>(); + + public static Metrics instance() { + return INSTANCE; + } + + public synchronized void setup(Meter meter, MetricsConfig metricsConfig) { + this.meter = meter; + this.globalConfig = metricsConfig; + setup0(); + } + + public HistogramBundle histogram(String name, String desc, String unit) { + return new HistogramBundle(name, desc, unit); + } + + public LongCounter counter(Function newFunc) { + return new LazyLongCounter(newFunc); + } + + public LongGaugeBundle longGauge(String name, String desc, String unit) { + return new LongGaugeBundle(name, desc, unit); + } + + public DoubleGaugeBundle doubleGauge(String name, String desc, String unit) { + return new DoubleGaugeBundle(name, desc, unit); + } + + private synchronized void setup0() { + if (meter == null) { + return; + } + for (; ; ) { + Setup setup = waitingSetups.poll(); + if (setup == null) { + break; + } + setup.setup(); + } + } + + interface Setup { + void setup(); + } + + public class HistogramBundle implements Setup { + private final List histograms = new CopyOnWriteArrayList<>(); + private final String name; + private final String desc; + private final String unit; + + private ObservableLongGauge count; + private ObservableLongGauge sum; + private ObservableDoubleGauge histP50Value; + private ObservableDoubleGauge histP99Value; + private ObservableDoubleGauge histMaxValue; + + @SuppressWarnings("this-escape") + public HistogramBundle(String name, String desc, String unit) { + this.name = name; + this.desc = desc; + this.unit = unit; + waitingSetups.add(this); + setup0(); + } + + public synchronized DeltaHistogram histogram(MetricsLevel level, Attributes attributes) { + Histogram histogram = new Histogram(level, attributes); + histograms.add(histogram); + histogram.setup(); + return histogram.histogram; + } + + public synchronized void setup() { + histograms.forEach(Histogram::setup); + this.count = meter.gaugeBuilder(name + S3StreamMetricsConstant.COUNT_METRIC_NAME_SUFFIX) + .setDescription(desc + " (count)") + .ofLongs() + .buildWithCallback(result -> { + histograms.forEach(histogram -> { + if (histogram.shouldRecord()) { + result.record(histogram.histogram.count(), histogram.attributes()); + } + }); + }); + this.sum = meter.gaugeBuilder(name + S3StreamMetricsConstant.SUM_METRIC_NAME_SUFFIX) + .setDescription(desc + " (sum)") + .ofLongs() + .setUnit(unit) + .buildWithCallback(result -> { + histograms.forEach(histogram -> { + if (histogram.shouldRecord()) { + result.record(histogram.histogram.sum(), histogram.attributes()); + } + }); + }); + this.histP50Value = meter.gaugeBuilder(name + S3StreamMetricsConstant.P50_METRIC_NAME_SUFFIX) + .setDescription(desc + " (50th percentile)") + .setUnit(unit) + .buildWithCallback(result -> { + histograms.forEach(histogram -> { + if (histogram.shouldRecord()) { + result.record(histogram.histogram.p50(), histogram.attributes()); + } + }); + }); + this.histP99Value = meter.gaugeBuilder(name + S3StreamMetricsConstant.P99_METRIC_NAME_SUFFIX) + .setDescription(desc + " (99th percentile)") + .setUnit(unit) + .buildWithCallback(result -> { + histograms.forEach(histogram -> { + if (histogram.shouldRecord()) { + result.record(histogram.histogram.p99(), histogram.attributes()); + } + }); + }); + this.histMaxValue = meter.gaugeBuilder(name + S3StreamMetricsConstant.MAX_METRIC_NAME_SUFFIX) + .setDescription(desc + " (max)") + .setUnit(unit) + .buildWithCallback(result -> { + histograms.forEach(histogram -> { + if (histogram.shouldRecord()) { + result.record(histogram.histogram.max(), histogram.attributes()); + } + }); + }); + } + + class Histogram { + final DeltaHistogram histogram; + final MetricsLevel level; + final Attributes histogramAttrs; + Attributes finalAttributes; + boolean shouldRecord = true; + + public Histogram(MetricsLevel level, Attributes attributes) { + this.histogram = new DeltaHistogram(); + this.level = level; + this.histogramAttrs = attributes; + this.finalAttributes = attributes; + } + + public Attributes attributes() { + return finalAttributes; + } + + public boolean shouldRecord() { + return shouldRecord; + } + + public void setup() { + if (meter == null) { + return; + } + this.finalAttributes = Attributes.builder() + .putAll(globalConfig.getBaseAttributes()) + .putAll(histogramAttrs) + .build(); + this.shouldRecord = level.isWithin(globalConfig.getMetricsLevel()); + histogram.setSnapshotInterval(globalConfig.getMetricsReportIntervalMs()); + } + } + } + + public class LongGaugeBundle implements Setup { + private final List gauges = new CopyOnWriteArrayList<>(); + private final String name; + private final String desc; + private final String unit; + + private ObservableLongGauge instrument; + + @SuppressWarnings("this-escape") + public LongGaugeBundle(String name, String desc, String unit) { + this.name = name; + this.desc = desc; + this.unit = unit; + waitingSetups.add(this); + setup0(); + } + + public LongGauge register(MetricsLevel level, Attributes attributes) { + LongGauge gauge = new LongGauge(level, attributes); + gauges.add(gauge); + gauge.setup(); + return gauge; + } + + public synchronized void setup() { + gauges.forEach(LongGauge::setup); + this.instrument = meter.gaugeBuilder(name) + .setDescription(desc) + .setUnit(unit) + .ofLongs() + .buildWithCallback(measurement -> gauges.forEach(gauge -> gauge.record(measurement))); + } + + public final class LongGauge implements AutoCloseable { + private final MetricsLevel level; + private final Attributes gaugeAttributes; + private final AtomicLong value = new AtomicLong(); + private final AtomicBoolean hasValue = new AtomicBoolean(false); + private Attributes finalAttributes = Attributes.empty(); + private volatile boolean shouldRecord = true; + + private LongGauge(MetricsLevel level, Attributes attributes) { + this.level = level; + this.gaugeAttributes = attributes; + this.finalAttributes = attributes; + } + + private void setup() { + if (meter != null && globalConfig != null) { + this.finalAttributes = Attributes.builder() + .putAll(globalConfig.getBaseAttributes()) + .putAll(gaugeAttributes) + .build(); + this.shouldRecord = level.isWithin(globalConfig.getMetricsLevel()); + } else { + this.finalAttributes = gaugeAttributes; + this.shouldRecord = true; + } + } + + public void record(long newValue) { + value.set(newValue); + hasValue.set(true); + } + + public void clear() { + hasValue.set(false); + } + + private void record(ObservableLongMeasurement measurement) { + if (shouldRecord && hasValue.get()) { + measurement.record(value.get(), finalAttributes); + } + } + + @Override + public void close() { + gauges.remove(this); + hasValue.set(false); + } + } + } + + public class DoubleGaugeBundle implements Setup { + private final List gauges = new CopyOnWriteArrayList<>(); + private final String name; + private final String desc; + private final String unit; + + private ObservableDoubleGauge instrument; + + @SuppressWarnings("this-escape") + public DoubleGaugeBundle(String name, String desc, String unit) { + this.name = name; + this.desc = desc; + this.unit = unit; + waitingSetups.add(this); + setup0(); + } + + public DoubleGauge register(MetricsLevel level, Attributes attributes) { + DoubleGauge gauge = new DoubleGauge(level, attributes); + gauges.add(gauge); + gauge.setup(); + return gauge; + } + + public synchronized void setup() { + gauges.forEach(DoubleGauge::setup); + this.instrument = meter.gaugeBuilder(name) + .setDescription(desc) + .setUnit(unit) + .buildWithCallback(measurement -> gauges.forEach(gauge -> gauge.record(measurement))); + } + + public final class DoubleGauge implements AutoCloseable { + private final MetricsLevel level; + private final Attributes gaugeAttributes; + private final AtomicReference value = new AtomicReference<>(0.0); + private final AtomicBoolean hasValue = new AtomicBoolean(false); + private Attributes finalAttributes = Attributes.empty(); + private volatile boolean shouldRecord = true; + + private DoubleGauge(MetricsLevel level, Attributes attributes) { + this.level = level; + this.gaugeAttributes = attributes; + this.finalAttributes = attributes; + } + + private void setup() { + if (meter != null && globalConfig != null) { + this.finalAttributes = Attributes.builder() + .putAll(globalConfig.getBaseAttributes()) + .putAll(gaugeAttributes) + .build(); + this.shouldRecord = level.isWithin(globalConfig.getMetricsLevel()); + } else { + this.finalAttributes = gaugeAttributes; + this.shouldRecord = true; + } + } + + public void record(double newValue) { + value.set(newValue); + hasValue.set(true); + } + + public void clear() { + hasValue.set(false); + } + + private void record(ObservableDoubleMeasurement measurement) { + if (shouldRecord && hasValue.get()) { + measurement.record(value.get(), finalAttributes); + } + } + + @Override + public void close() { + gauges.remove(this); + hasValue.set(false); + } + } + } + + public class LazyLongCounter implements Setup, LongCounter { + private final Function newFunc; + private LongCounter counter = new NoopLongCounter(); + + @SuppressWarnings("this-escape") + public LazyLongCounter(Function newFunc) { + this.newFunc = newFunc; + waitingSetups.add(this); + setup0(); + } + + @Override + public void setup() { + this.counter = newFunc.apply(meter); + } + + @Override + public void add(long value) { + counter.add(value); + } + + @Override + public void add(long value, Attributes attributes) { + counter.add(value, attributes); + } + + @Override + public void add(long value, Attributes attributes, Context context) { + counter.add(value, attributes, context); + } + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsConfig.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsConfig.java index 9ab79bf740..5eb07b0965 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsConfig.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsConfig.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsLevel.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsLevel.java index 36332fd664..e43791d8ca 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsLevel.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/MetricsLevel.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongCounter.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongCounter.java index be3854f130..72b3a7ca59 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongCounter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongCounter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongHistogram.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongHistogram.java index df5b47f868..7411656131 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongHistogram.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopLongHistogram.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableDoubleGauge.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableDoubleGauge.java new file mode 100644 index 0000000000..880ab9056e --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableDoubleGauge.java @@ -0,0 +1,25 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.metrics; + +import io.opentelemetry.api.metrics.ObservableDoubleGauge; + +public class NoopObservableDoubleGauge implements ObservableDoubleGauge { +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableLongGauge.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableLongGauge.java index 3945d122f6..093fb0c50e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableLongGauge.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/NoopObservableLongGauge.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsConstant.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsConstant.java index 38f07a4117..0de214b76f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsConstant.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsConstant.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; @@ -151,4 +159,8 @@ public class S3StreamMetricsConstant { public static final String LABEL_STAGE_GET_OBJECTS = "get_objects"; public static final String LABEL_STAGE_FIND_INDEX = "find_index"; public static final String LABEL_STAGE_COMPUTE = "compute"; + + // Broker Quota + public static final String BROKER_QUOTA_LIMIT_METRIC_NAME = "broker_quota_limit"; + public static final AttributeKey LABEL_BROKER_QUOTA_TYPE = AttributeKey.stringKey("type"); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsManager.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsManager.java index 3653873bd8..436b67a909 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsManager.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/S3StreamMetricsManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; @@ -34,6 +42,7 @@ import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.metrics.LongCounter; import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.metrics.ObservableDoubleGauge; import io.opentelemetry.api.metrics.ObservableLongGauge; import static com.automq.stream.s3.metrics.S3StreamMetricsConstant.LABEL_CACHE_NAME; @@ -139,10 +148,16 @@ public class S3StreamMetricsManager { private static final MultiAttributes OPERATOR_INDEX_ATTRIBUTES = new MultiAttributes<>(Attributes.empty(), S3StreamMetricsConstant.LABEL_INDEX); + // Broker Quota + private static final MultiAttributes BROKER_QUOTA_TYPE_ATTRIBUTES = new MultiAttributes<>(Attributes.empty(), + S3StreamMetricsConstant.LABEL_BROKER_QUOTA_TYPE); + private static ObservableDoubleGauge brokerQuotaLimit = new NoopObservableDoubleGauge(); + private static Supplier> brokerQuotaLimitSupplier = () -> new ConcurrentHashMap<>(); static { BASE_ATTRIBUTES_LISTENERS.add(ALLOC_TYPE_ATTRIBUTES); BASE_ATTRIBUTES_LISTENERS.add(OPERATOR_INDEX_ATTRIBUTES); + BASE_ATTRIBUTES_LISTENERS.add(BROKER_QUOTA_TYPE_ATTRIBUTES); } public static void configure(MetricsConfig metricsConfig) { @@ -400,6 +415,7 @@ public static void initMetrics(Meter meter, String prefix) { }); initAsyncCacheMetrics(meter, prefix); + initBrokerQuotaMetrics(meter, prefix); } private static void initAsyncCacheMetrics(Meter meter, String prefix) { @@ -475,21 +491,48 @@ private static void initAsyncCacheMetrics(Meter meter, String prefix) { }); } - public static void registerNetworkLimiterSupplier(AsyncNetworkBandwidthLimiter.Type type, - Supplier networkAvailableBandwidthSupplier, + private static void initBrokerQuotaMetrics(Meter meter, String prefix) { + brokerQuotaLimit = meter.gaugeBuilder(prefix + S3StreamMetricsConstant.BROKER_QUOTA_LIMIT_METRIC_NAME) + .setDescription("Broker quota limit") + .buildWithCallback(result -> { + if (MetricsLevel.INFO.isWithin(metricsConfig.getMetricsLevel())) { + Map brokerQuotaLimitMap = brokerQuotaLimitSupplier.get(); + for (Map.Entry entry : brokerQuotaLimitMap.entrySet()) { + String quotaType = entry.getKey(); + Double quotaLimit = entry.getValue(); + // drop too large values + if (quotaLimit > 1e15) { + continue; + } + result.record(quotaLimit, BROKER_QUOTA_TYPE_ATTRIBUTES.get(quotaType)); + } + } + }); + } + + public static void registerNetworkLimiterQueueSizeSupplier(AsyncNetworkBandwidthLimiter.Type type, Supplier networkLimiterQueueSizeSupplier) { switch (type) { case INBOUND: - S3StreamMetricsManager.networkInboundAvailableBandwidthSupplier = networkAvailableBandwidthSupplier; S3StreamMetricsManager.networkInboundLimiterQueueSizeSupplier = networkLimiterQueueSizeSupplier; break; case OUTBOUND: - S3StreamMetricsManager.networkOutboundAvailableBandwidthSupplier = networkAvailableBandwidthSupplier; S3StreamMetricsManager.networkOutboundLimiterQueueSizeSupplier = networkLimiterQueueSizeSupplier; break; } } + public static void registerNetworkAvailableBandwidthSupplier(AsyncNetworkBandwidthLimiter.Type type, Supplier networkAvailableBandwidthSupplier) { + switch (type) { + case INBOUND: + S3StreamMetricsManager.networkInboundAvailableBandwidthSupplier = networkAvailableBandwidthSupplier; + break; + case OUTBOUND: + S3StreamMetricsManager.networkOutboundAvailableBandwidthSupplier = networkAvailableBandwidthSupplier; + break; + } + } + public static void registerDeltaWalOffsetSupplier(Supplier deltaWalStartOffsetSupplier, Supplier deltaWalTrimmedOffsetSupplier) { S3StreamMetricsManager.deltaWalStartOffsetSupplier = deltaWalStartOffsetSupplier; @@ -899,4 +942,8 @@ public static void registerLocalStreamRangeIndexCacheSizeSupplier(Supplier localStreamRangeIndexCacheStreamNum) { S3StreamMetricsManager.localStreamRangeIndexCacheStreamNum = localStreamRangeIndexCacheStreamNum; } + + public static void registerBrokerQuotaLimitSupplier(Supplier> brokerQuotaLimitSupplier) { + S3StreamMetricsManager.brokerQuotaLimitSupplier = brokerQuotaLimitSupplier; + } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/TimerUtil.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/TimerUtil.java index dd43d98053..0cb657d329 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/TimerUtil.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/TimerUtil.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3MetricsType.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3MetricsType.java index 4738704265..00abb821a7 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3MetricsType.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3MetricsType.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.operations; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3ObjectStage.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3ObjectStage.java index 89917671dd..d9495b73b0 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3ObjectStage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3ObjectStage.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.operations; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Operation.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Operation.java index 0745c89246..2249486807 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Operation.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Operation.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.operations; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Stage.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Stage.java index fdc58a8cac..c5be1496e1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Stage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3Stage.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.operations; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/AsyncLRUCacheStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/AsyncLRUCacheStats.java index 07c5fcb234..7e520f9a1f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/AsyncLRUCacheStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/AsyncLRUCacheStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; @@ -36,7 +44,7 @@ private AsyncLRUCacheStats() { public static AsyncLRUCacheStats getInstance() { if (instance == null) { - synchronized (NetworkStats.class) { + synchronized (AsyncLRUCacheStats.class) { if (instance == null) { instance = new AsyncLRUCacheStats(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/CompactionStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/CompactionStats.java index fb42feef01..49cebd58b6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/CompactionStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/CompactionStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/MetadataStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/MetadataStats.java index 7fe6638dc3..cb132caadf 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/MetadataStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/MetadataStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/NetworkStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/NetworkStats.java index 572c80387e..19c179141f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/NetworkStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/NetworkStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; @@ -27,7 +35,7 @@ import java.util.concurrent.ConcurrentHashMap; public class NetworkStats { - private static volatile NetworkStats instance = null; + private static final NetworkStats INSTANCE = new NetworkStats(); // > private final Map> streamReadBytesStats = new ConcurrentHashMap<>(); private final Counter networkInboundUsageTotal = new Counter(); @@ -41,20 +49,20 @@ private NetworkStats() { } public static NetworkStats getInstance() { - if (instance == null) { - synchronized (NetworkStats.class) { - if (instance == null) { - instance = new NetworkStats(); - } - } - } - return instance; + return INSTANCE; } public CounterMetric networkUsageTotalStats(AsyncNetworkBandwidthLimiter.Type type, ThrottleStrategy strategy) { - return type == AsyncNetworkBandwidthLimiter.Type.INBOUND - ? networkInboundUsageTotalStats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkInboundUsageMetric(strategy, networkInboundUsageTotal::inc)) - : networkOutboundUsageTotalStats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkOutboundUsageMetric(strategy, networkOutboundUsageTotal::inc)); + Map stats = type == AsyncNetworkBandwidthLimiter.Type.INBOUND ? networkInboundUsageTotalStats : networkOutboundUsageTotalStats; + CounterMetric metric = stats.get(strategy); + if (metric == null) { + if (type == AsyncNetworkBandwidthLimiter.Type.INBOUND) { + metric = stats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkInboundUsageMetric(strategy, networkInboundUsageTotal::inc)); + } else { + metric = stats.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkOutboundUsageMetric(strategy, networkOutboundUsageTotal::inc)); + } + } + return metric; } public Optional fastReadBytesStats(long streamId) { @@ -88,8 +96,18 @@ public Map> allStreamReadBytesStats() { } public HistogramMetric networkLimiterQueueTimeStats(AsyncNetworkBandwidthLimiter.Type type, ThrottleStrategy strategy) { - return type == AsyncNetworkBandwidthLimiter.Type.INBOUND - ? networkInboundLimiterQueueTimeStatsMap.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkInboundLimiterQueueTimeMetric(MetricsLevel.INFO, strategy)) - : networkOutboundLimiterQueueTimeStatsMap.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkOutboundLimiterQueueTimeMetric(MetricsLevel.INFO, strategy)); + HistogramMetric metric; + if (type == AsyncNetworkBandwidthLimiter.Type.INBOUND) { + metric = networkInboundLimiterQueueTimeStatsMap.get(strategy); + if (metric == null) { + metric = networkInboundLimiterQueueTimeStatsMap.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkInboundLimiterQueueTimeMetric(MetricsLevel.INFO, strategy)); + } + } else { + metric = networkOutboundLimiterQueueTimeStatsMap.get(strategy); + if (metric == null) { + metric = networkOutboundLimiterQueueTimeStatsMap.computeIfAbsent(strategy, k -> S3StreamMetricsManager.buildNetworkOutboundLimiterQueueTimeMetric(MetricsLevel.INFO, strategy)); + } + } + return metric; } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3ObjectStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3ObjectStats.java index de324a5629..e856f8270e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3ObjectStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3ObjectStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3OperationStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3OperationStats.java index 6f01f742e5..6ef5c6334e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3OperationStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/S3OperationStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; @@ -66,7 +74,7 @@ private S3OperationStats() { public static S3OperationStats getInstance() { if (instance == null) { - synchronized (StreamOperationStats.class) { + synchronized (S3OperationStats.class) { if (instance == null) { instance = new S3OperationStats(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StorageOperationStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StorageOperationStats.java index 1db9cfe703..c8cb6cddb5 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StorageOperationStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StorageOperationStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StreamOperationStats.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StreamOperationStats.java index 8a3530b87b..cd458187ba 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StreamOperationStats.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/StreamOperationStats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.stats; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigListener.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigListener.java index 3f61d4529d..67c21354c0 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigListener.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigListener.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.wrapper; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigurableMetric.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigurableMetric.java index fd3ef33c48..56c2b80469 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigurableMetric.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/ConfigurableMetric.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.wrapper; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/Counter.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/Counter.java index 98e38c309a..cbadffa9cc 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/Counter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/Counter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.wrapper; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/CounterMetric.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/CounterMetric.java index 216b567db6..a4d1f6d3db 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/CounterMetric.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/CounterMetric.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.wrapper; diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/DeltaHistogram.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/DeltaHistogram.java index 9c6421ca18..6b7ed7b290 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/DeltaHistogram.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/DeltaHistogram.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.wrapper; @@ -14,12 +22,13 @@ import org.HdrHistogram.Histogram; import org.HdrHistogram.Recorder; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.LongAdder; import java.util.function.BiPredicate; public class DeltaHistogram { - private static final Long DEFAULT_SNAPSHOT_INTERVAL_MS = 5000L; + private static final Long DEFAULT_SNAPSHOT_INTERVAL_MS = TimeUnit.SECONDS.toMillis(30); private final LongAdder cumulativeCount = new LongAdder(); private final LongAdder cumulativeSum = new LongAdder(); private final AtomicLong min = new AtomicLong(Long.MAX_VALUE); @@ -124,36 +133,49 @@ public double p50() { return lastSnapshot.p50; } - private void snapshotAndReset() { + public SnapshotExt snapshotAndReset() { synchronized (this) { if (lastSnapshot == null || System.currentTimeMillis() - lastSnapshotTime > snapshotInterval) { - this.intervalHistogram = this.recorder.getIntervalHistogram(this.intervalHistogram); + snapshotAndReset0(); + } + } + return lastSnapshot; + } - long snapshotMin = min.get(); - long snapshotMax = max.get(); + private void snapshotAndReset0() { + this.intervalHistogram = this.recorder.getIntervalHistogram(this.intervalHistogram); - long newCount = cumulativeCount.sum(); - long newSum = cumulativeSum.sum(); + long snapshotMin = min.get(); + long snapshotMax = max.get(); - long snapshotCount = newCount - lastCount; - long snapshotSum = newSum - lastSum; + long newCount = cumulativeCount.sum(); + long newSum = cumulativeSum.sum(); - double p99 = intervalHistogram.getValueAtPercentile(0.99); - double p95 = intervalHistogram.getValueAtPercentile(0.95); - double p50 = intervalHistogram.getValueAtPercentile(0.50); + long snapshotCount = newCount - lastCount; + long snapshotSum = newSum - lastSum; - lastCount = newCount; - lastSum = newSum; + double p99 = intervalHistogram.getValueAtPercentile(99); + double p95 = intervalHistogram.getValueAtPercentile(95); + double p50 = intervalHistogram.getValueAtPercentile(50); - min.set(0); - max.set(0); - lastSnapshot = new SnapshotExt(snapshotMin, snapshotMax, snapshotCount, snapshotSum, p99, p95, p50); - lastSnapshotTime = System.currentTimeMillis(); - } + lastCount = newCount; + lastSum = newSum; + + min.set(0); + max.set(0); + lastSnapshot = new SnapshotExt(snapshotMin, snapshotMax, snapshotCount, snapshotSum, p99, p95, p50); + lastSnapshotTime = System.currentTimeMillis(); + } + + // for benchmark only + public SnapshotExt forceSnapshotAndReset() { + synchronized (this) { + snapshotAndReset0(); } + return lastSnapshot; } - static class SnapshotExt { + public static class SnapshotExt { final long min; final long max; final long count; @@ -172,12 +194,40 @@ public SnapshotExt(long min, long max, long count, long sum, double p99, double this.p95 = p95; } - double mean() { + public double mean() { if (count == 0) { return 0; } else { return (double) sum / count; } } + + public long getMin() { + return min; + } + + public long getMax() { + return max; + } + + public long getCount() { + return count; + } + + public long getSum() { + return sum; + } + + public double getP99() { + return p99; + } + + public double getP95() { + return p95; + } + + public double getP50() { + return p50; + } } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/HistogramMetric.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/HistogramMetric.java index 313b123dbe..ad86ed00e1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/HistogramMetric.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/HistogramMetric.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics.wrapper; @@ -27,6 +35,7 @@ public HistogramMetric(MetricsLevel currentMetricsLevel, MetricsConfig metricsCo public HistogramMetric(MetricsLevel currentMetricsLevel, MetricsConfig metricsConfig, Attributes extraAttributes) { super(metricsConfig, extraAttributes); this.deltaHistogram = new DeltaHistogram(); + this.deltaHistogram.setSnapshotInterval(metricsConfig.getMetricsReportIntervalMs()); this.currentMetricsLevel = currentMetricsLevel; } diff --git a/s3stream/src/main/java/com/automq/stream/s3/model/StreamRecordBatch.java b/s3stream/src/main/java/com/automq/stream/s3/model/StreamRecordBatch.java index 3de2038f42..20d2bbe81f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/model/StreamRecordBatch.java +++ b/s3stream/src/main/java/com/automq/stream/s3/model/StreamRecordBatch.java @@ -1,55 +1,70 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.model; -import com.automq.stream.s3.StreamRecordBatchCodec; +import com.automq.stream.ByteBufSeqAlloc; import com.automq.stream.utils.biniarysearch.ComparableItem; +import java.nio.ByteBuffer; + import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +import static com.automq.stream.s3.ByteBufAlloc.DECODE_RECORD; +import static com.automq.stream.s3.ByteBufAlloc.ENCODE_RECORD; +import static com.automq.stream.s3.StreamRecordBatchCodec.BASE_OFFSET_POS; +import static com.automq.stream.s3.StreamRecordBatchCodec.EPOCH_POS; +import static com.automq.stream.s3.StreamRecordBatchCodec.HEADER_SIZE; +import static com.automq.stream.s3.StreamRecordBatchCodec.LAST_OFFSET_DELTA_POS; +import static com.automq.stream.s3.StreamRecordBatchCodec.MAGIC_POS; +import static com.automq.stream.s3.StreamRecordBatchCodec.MAGIC_V0; +import static com.automq.stream.s3.StreamRecordBatchCodec.PAYLOAD_LENGTH_POS; +import static com.automq.stream.s3.StreamRecordBatchCodec.PAYLOAD_POS; +import static com.automq.stream.s3.StreamRecordBatchCodec.STREAM_ID_POS; public class StreamRecordBatch implements Comparable, ComparableItem { private static final int OBJECT_OVERHEAD = 48 /* fields */ + 48 /* ByteBuf payload */ + 48 /* ByteBuf encoded */; - private final long streamId; - private final long epoch; + private static final ByteBufSeqAlloc ENCODE_ALLOC = new ByteBufSeqAlloc(ENCODE_RECORD, 8); + private static final ByteBufSeqAlloc DECODE_ALLOC = new ByteBufSeqAlloc(DECODE_RECORD, 8); + // Cache the frequently used fields private final long baseOffset; private final int count; - private ByteBuf payload; - private ByteBuf encoded; - public StreamRecordBatch(long streamId, long epoch, long baseOffset, int count, ByteBuf payload) { - this.streamId = streamId; - this.epoch = epoch; - this.baseOffset = baseOffset; - this.count = count; - this.payload = payload; + final ByteBuf encoded; + + private StreamRecordBatch(ByteBuf encoded) { + this.encoded = encoded; + this.baseOffset = encoded.getLong(encoded.readerIndex() + BASE_OFFSET_POS); + this.count = encoded.getInt(encoded.readerIndex() + LAST_OFFSET_DELTA_POS); } public ByteBuf encoded() { - // TODO: keep the ref count - if (encoded == null) { - encoded = StreamRecordBatchCodec.encode(this); - ByteBuf oldPayload = payload; - payload = encoded.slice(encoded.readerIndex() + encoded.readableBytes() - payload.readableBytes(), payload.readableBytes()); - oldPayload.release(); - } - return encoded.duplicate(); + return encoded.slice(); } public long getStreamId() { - return streamId; + return encoded.getLong(encoded.readerIndex() + STREAM_ID_POS); } public long getEpoch() { - return epoch; + return encoded.getLong(encoded.readerIndex() + EPOCH_POS); } public long getBaseOffset() { @@ -57,7 +72,14 @@ public long getBaseOffset() { } public long getLastOffset() { - return baseOffset + count; + long baseOffset = getBaseOffset(); + int count = getCount(); + if (count > 0) { + return baseOffset + count; + } else { + // link record + return baseOffset - count; + } } public int getCount() { @@ -65,11 +87,11 @@ public int getCount() { } public ByteBuf getPayload() { - return payload; + return encoded.slice(encoded.readerIndex() + PAYLOAD_POS, encoded.readableBytes() - HEADER_SIZE); } public int size() { - return payload.readableBytes(); + return encoded.getInt(encoded.readerIndex() + PAYLOAD_LENGTH_POS); } public int occupiedSize() { @@ -77,41 +99,29 @@ public int occupiedSize() { } public void retain() { - if (encoded != null) { - encoded.retain(); - } else { - payload.retain(); - } + encoded.retain(); } public void release() { - if (encoded != null) { - encoded.release(); - } else { - payload.release(); - } + encoded.release(); } @Override public int compareTo(StreamRecordBatch o) { - int rst = Long.compare(streamId, o.streamId); + int rst = Long.compare(getStreamId(), o.getStreamId()); if (rst != 0) { return rst; } - rst = Long.compare(epoch, o.epoch); - if (rst != 0) { - return rst; - } - return Long.compare(baseOffset, o.baseOffset); + return Long.compare(getBaseOffset(), o.getBaseOffset()); } @Override public String toString() { return "StreamRecordBatch{" + - "streamId=" + streamId + - ", epoch=" + epoch + - ", baseOffset=" + baseOffset + - ", count=" + count + + "streamId=" + getStreamId() + + ", epoch=" + getEpoch() + + ", baseOffset=" + getBaseOffset() + + ", count=" + getCount() + ", size=" + size() + '}'; } @@ -124,4 +134,67 @@ public boolean isLessThan(Long value) { public boolean isGreaterThan(Long value) { return getBaseOffset() > value; } + + public static StreamRecordBatch of(long streamId, long epoch, long baseOffset, int count, ByteBuffer payload) { + return of(streamId, epoch, baseOffset, count, Unpooled.wrappedBuffer(payload), ENCODE_ALLOC); + } + + public static StreamRecordBatch of(long streamId, long epoch, long baseOffset, int count, ByteBuffer payload, ByteBufSeqAlloc alloc) { + return of(streamId, epoch, baseOffset, count, Unpooled.wrappedBuffer(payload), alloc); + } + + /** + * StreamRecordBatch.of expects take the owner of the payload. + * The payload will be copied to the new StreamRecordBatch and released. + */ + public static StreamRecordBatch of(long streamId, long epoch, long baseOffset, int count, ByteBuf payload) { + return of(streamId, epoch, baseOffset, count, payload, ENCODE_ALLOC); + } + + /** + * StreamRecordBatch.of expects take the owner of the payload. + * The payload will be copied to the new StreamRecordBatch and released. + */ + public static StreamRecordBatch of(long streamId, long epoch, long baseOffset, int count, ByteBuf payload, + ByteBufSeqAlloc alloc) { + int totalLength = HEADER_SIZE + payload.readableBytes(); + ByteBuf buf = alloc.byteBuffer(totalLength); + buf.writeByte(MAGIC_V0); + buf.writeLong(streamId); + buf.writeLong(epoch); + buf.writeLong(baseOffset); + buf.writeInt(count); + buf.writeInt(payload.readableBytes()); + buf.writeBytes(payload); + payload.release(); + return new StreamRecordBatch(buf); + } + + public static StreamRecordBatch parse(ByteBuf buf, boolean duplicated) { + return parse(buf, duplicated, DECODE_ALLOC); + } + + /** + * Won't release the input ByteBuf. + * - If duplicated is true, the returned StreamRecordBatch has its own copy of the data. + * - If duplicated is false, the returned StreamRecordBatch shares and retains the data buffer with the input. + */ + public static StreamRecordBatch parse(ByteBuf buf, boolean duplicated, ByteBufSeqAlloc alloc) { + int readerIndex = buf.readerIndex(); + byte magic = buf.getByte(readerIndex + MAGIC_POS); + if (magic != MAGIC_V0) { + throw new RuntimeException("Invalid magic byte " + magic); + } + int payloadSize = buf.getInt(readerIndex + PAYLOAD_LENGTH_POS); + int encodedSize = PAYLOAD_POS + payloadSize; + if (duplicated) { + ByteBuf encoded = alloc.byteBuffer(encodedSize); + buf.readBytes(encoded, encodedSize); + return new StreamRecordBatch(encoded); + } else { + ByteBuf encoded = buf.retainedSlice(readerIndex, encodedSize); + buf.skipBytes(encodedSize); + return new StreamRecordBatch(encoded); + } + } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java b/s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java index 483be75e40..0eee0fd578 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.network; @@ -15,6 +23,7 @@ import com.automq.stream.s3.metrics.S3StreamMetricsManager; import com.automq.stream.s3.metrics.stats.NetworkStats; import com.automq.stream.utils.LogContext; +import com.automq.stream.utils.Threads; import org.slf4j.Logger; @@ -23,9 +32,9 @@ import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; @@ -43,7 +52,7 @@ public class AsyncNetworkBandwidthLimiter implements NetworkBandwidthLimiter { private final Queue queuedCallbacks; private final Type type; private final long tokenSize; - private long availableTokens; + private final AtomicLong availableTokens; public AsyncNetworkBandwidthLimiter(Type type, long tokenSize, int refillIntervalMs) { this(type, tokenSize, refillIntervalMs, tokenSize); @@ -53,14 +62,16 @@ public AsyncNetworkBandwidthLimiter(Type type, long tokenSize, int refillInterva public AsyncNetworkBandwidthLimiter(Type type, long tokenSize, int refillIntervalMs, long maxTokens) { this.type = type; this.tokenSize = tokenSize; - this.availableTokens = this.tokenSize; + this.availableTokens = new AtomicLong(this.tokenSize); this.maxTokens = maxTokens; this.queuedCallbacks = new PriorityQueue<>(); - this.refillThreadPool = Executors.newSingleThreadScheduledExecutor(new DefaultThreadFactory("refill-bucket-thread")); - this.callbackThreadPool = Executors.newFixedThreadPool(1, new DefaultThreadFactory("callback-thread")); + this.refillThreadPool = + Threads.newSingleThreadScheduledExecutor(new DefaultThreadFactory("refill-bucket-thread"), LOGGER); + // The threads number must be larger than 1 because the #run will occupy one thread. + this.callbackThreadPool = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(2, "callback-thread", true, LOGGER); this.callbackThreadPool.execute(this::run); this.refillThreadPool.scheduleAtFixedRate(this::refillToken, refillIntervalMs, refillIntervalMs, TimeUnit.MILLISECONDS); - S3StreamMetricsManager.registerNetworkLimiterSupplier(type, this::getAvailableTokens, this::getQueueSize); + S3StreamMetricsManager.registerNetworkLimiterQueueSizeSupplier(type, this::getQueueSize); LOGGER.info("AsyncNetworkBandwidthLimiter initialized, type: {}, tokenSize: {}, maxTokens: {}, refillIntervalMs: {}", type.getName(), tokenSize, maxTokens, refillIntervalMs); } @@ -79,7 +90,7 @@ private void run() { } long size = Math.min(head.size, MAX_TOKEN_PART_SIZE); reduceToken(size); - if (head.complete(size)) { + if (head.complete(size, callbackThreadPool)) { queuedCallbacks.poll(); } } @@ -94,7 +105,7 @@ private void run() { private void refillToken() { lock.lock(); try { - availableTokens = Math.min(availableTokens + this.tokenSize, this.maxTokens); + this.availableTokens.getAndUpdate(old -> Math.min(old + this.tokenSize, this.maxTokens)); condition.signalAll(); } finally { lock.unlock(); @@ -105,7 +116,7 @@ private boolean ableToConsume() { if (queuedCallbacks.isEmpty()) { return false; } - return availableTokens > 0; + return availableTokens.get() > 0; } public void shutdown() { @@ -118,12 +129,7 @@ public long getMaxTokens() { } public long getAvailableTokens() { - lock.lock(); - try { - return availableTokens; - } finally { - lock.unlock(); - } + return availableTokens.get(); } public int getQueueSize() { @@ -136,12 +142,7 @@ public int getQueueSize() { } private void forceConsume(long size) { - lock.lock(); - try { - reduceToken(size); - } finally { - lock.unlock(); - } + reduceToken(size); } public CompletableFuture consume(ThrottleStrategy throttleStrategy, long size) { @@ -151,24 +152,28 @@ public CompletableFuture consume(ThrottleStrategy throttleStrategy, long s forceConsume(size); cf.complete(null); } else { + boolean satisfied = false; lock.lock(); try { - if (availableTokens <= 0 || !queuedCallbacks.isEmpty()) { + if (availableTokens.get() <= 0 || !queuedCallbacks.isEmpty()) { queuedCallbacks.offer(new BucketItem(throttleStrategy, size, cf)); condition.signalAll(); } else { reduceToken(size); - cf.complete(null); + satisfied = true; } } finally { lock.unlock(); } + if (satisfied) { + cf.complete(null); + } } return cf; } private void reduceToken(long size) { - this.availableTokens = Math.max(-maxTokens, availableTokens - size); + this.availableTokens.getAndUpdate(old -> Math.max(-maxTokens, old - size)); } public enum Type { @@ -186,7 +191,7 @@ public String getName() { } } - private static class BucketItem implements Comparable { + private class BucketItem implements Comparable { private final ThrottleStrategy strategy; private final CompletableFuture cf; private final long timestamp; @@ -207,10 +212,11 @@ public int compareTo(BucketItem o) { return Long.compare(strategy.priority(), o.strategy.priority()); } - public boolean complete(long completeSize) { + public boolean complete(long completeSize, ExecutorService executor) { size -= completeSize; if (size <= 0) { - cf.complete(null); + executor.submit(() -> cf.complete(null)); + NetworkStats.getInstance().networkLimiterQueueTimeStats(type, strategy).record(System.nanoTime() - timestamp); return true; } return false; diff --git a/s3stream/src/main/java/com/automq/stream/s3/network/GlobalNetworkBandwidthLimiters.java b/s3stream/src/main/java/com/automq/stream/s3/network/GlobalNetworkBandwidthLimiters.java index 69ea729e65..3a1961853f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/network/GlobalNetworkBandwidthLimiters.java +++ b/s3stream/src/main/java/com/automq/stream/s3/network/GlobalNetworkBandwidthLimiters.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.network; @@ -18,6 +26,8 @@ public class GlobalNetworkBandwidthLimiters { private final Map limiters = new HashMap<>(); private static final GlobalNetworkBandwidthLimiters INSTANCE = new GlobalNetworkBandwidthLimiters(); + private NetworkBandwidthLimiter inboundLimiter = AsyncNetworkBandwidthLimiter.NOOP; + private NetworkBandwidthLimiter outboundLimiter = AsyncNetworkBandwidthLimiter.NOOP; public static GlobalNetworkBandwidthLimiters instance() { return INSTANCE; @@ -28,6 +38,11 @@ public void setup(AsyncNetworkBandwidthLimiter.Type type, long tokenSize, int re throw new IllegalArgumentException(type + " is already setup"); } limiters.put(type, new AsyncNetworkBandwidthLimiter(type, tokenSize, refillIntervalMs, maxTokens)); + if (type == AsyncNetworkBandwidthLimiter.Type.INBOUND) { + inboundLimiter = limiters.get(type); + } else if (type == AsyncNetworkBandwidthLimiter.Type.OUTBOUND) { + outboundLimiter = limiters.get(type); + } } public NetworkBandwidthLimiter get(AsyncNetworkBandwidthLimiter.Type type) { @@ -38,4 +53,12 @@ public NetworkBandwidthLimiter get(AsyncNetworkBandwidthLimiter.Type type) { return limiter; } + public NetworkBandwidthLimiter inbound() { + return inboundLimiter; + } + + public NetworkBandwidthLimiter outbound() { + return outboundLimiter; + } + } diff --git a/s3stream/src/main/java/com/automq/stream/s3/network/NetworkBandwidthLimiter.java b/s3stream/src/main/java/com/automq/stream/s3/network/NetworkBandwidthLimiter.java index 69b32b8d15..e6e957c2c6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/network/NetworkBandwidthLimiter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/network/NetworkBandwidthLimiter.java @@ -1,23 +1,46 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.network; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; public interface NetworkBandwidthLimiter { NetworkBandwidthLimiter NOOP = new Noop(); CompletableFuture consume(ThrottleStrategy throttleStrategy, long size); + default void consumeBlocking(ThrottleStrategy throttleStrategy, long size) + throws InterruptedException, ExecutionException { + CompletableFuture future = consume(throttleStrategy, size); + if (future == null) { + return; + } + try { + future.get(); + } catch (InterruptedException e) { + future.cancel(true); + throw e; + } + } + long getMaxTokens(); long getAvailableTokens(); diff --git a/s3stream/src/main/java/com/automq/stream/s3/network/ThrottleStrategy.java b/s3stream/src/main/java/com/automq/stream/s3/network/ThrottleStrategy.java index 17c7a52e85..0ff361645e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/network/ThrottleStrategy.java +++ b/s3stream/src/main/java/com/automq/stream/s3/network/ThrottleStrategy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.network; @@ -15,7 +23,8 @@ public enum ThrottleStrategy { BYPASS(0, "bypass"), COMPACTION(1, "compaction"), TAIL(2, "tail"), - CATCH_UP(3, "catchup"); + CATCH_UP(3, "catchup"), + ICEBERG_WRITE(4, "iceberg_write"); private final int priority; private final String name; diff --git a/s3stream/src/main/java/com/automq/stream/s3/network/test/RecordTestNetworkBandwidthLimiter.java b/s3stream/src/main/java/com/automq/stream/s3/network/test/RecordTestNetworkBandwidthLimiter.java index b7e0042f42..8c05691913 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/network/test/RecordTestNetworkBandwidthLimiter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/network/test/RecordTestNetworkBandwidthLimiter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.network.test; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectHook.java b/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectHook.java index e2071b3a48..803288a289 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectHook.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectHook.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectRequest.java b/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectRequest.java index 496e488e97..5e93256934 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectRequest.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectResponse.java b/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectResponse.java index 218d99dffc..6094f60bd8 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectResponse.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/CommitStreamSetObjectResponse.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/CompactStreamObjectRequest.java b/s3stream/src/main/java/com/automq/stream/s3/objects/CompactStreamObjectRequest.java index 5f6c675573..55cd1facbf 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/CompactStreamObjectRequest.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/CompactStreamObjectRequest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectAttributes.java b/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectAttributes.java index bd05d034f8..2df0fb1516 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectAttributes.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectAttributes.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectManager.java b/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectManager.java index 5aae0d52a6..203c79f160 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectManager.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectStreamRange.java b/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectStreamRange.java index 445651ada4..3e601d2bcf 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectStreamRange.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/ObjectStreamRange.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; @@ -26,6 +34,7 @@ public ObjectStreamRange(long streamId, long epoch, long startOffset, long endOf this.epoch = epoch; this.startOffset = startOffset; this.endOffset = endOffset; + // TODO: remove useless size this.size = size; } diff --git a/s3stream/src/main/java/com/automq/stream/s3/objects/StreamObject.java b/s3stream/src/main/java/com/automq/stream/s3/objects/StreamObject.java index d01f7f1e00..28351d1265 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/objects/StreamObject.java +++ b/s3stream/src/main/java/com/automq/stream/s3/objects/StreamObject.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java index 1d757b1ec3..11d98eac22 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -16,24 +24,23 @@ import com.automq.stream.s3.metrics.S3StreamMetricsManager; import com.automq.stream.s3.metrics.TimerUtil; import com.automq.stream.s3.metrics.operations.S3Operation; -import com.automq.stream.s3.metrics.stats.NetworkStats; import com.automq.stream.s3.metrics.stats.S3OperationStats; import com.automq.stream.s3.metrics.stats.StorageOperationStats; -import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter; import com.automq.stream.s3.network.NetworkBandwidthLimiter; import com.automq.stream.s3.network.ThrottleStrategy; import com.automq.stream.s3.objects.ObjectAttributes; import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.LogContext; import com.automq.stream.utils.ThreadUtils; import com.automq.stream.utils.Threads; import com.automq.stream.utils.Utils; import org.apache.commons.lang3.tuple.Pair; import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import java.time.Duration; import java.util.ArrayList; +import java.util.Comparator; import java.util.HashMap; import java.util.Iterator; import java.util.LinkedList; @@ -41,32 +48,40 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; +import java.util.concurrent.PriorityBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.Semaphore; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.BiFunction; +import java.util.function.Supplier; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.util.HashedWheelTimer; import io.netty.util.ReferenceCounted; +import software.amazon.awssdk.http.HttpStatusCode; +import software.amazon.awssdk.services.s3.model.S3Exception; @SuppressWarnings("this-escape") public abstract class AbstractObjectStorage implements ObjectStorage { - static final Logger LOGGER = LoggerFactory.getLogger(AbstractObjectStorage.class); private static final int MAX_INFLIGHT_FAST_RETRY_COUNT = 5; - private static final int DEFAULT_RETRY_DELAY = 100; private static final AtomicInteger INDEX = new AtomicInteger(-1); private static final int DEFAULT_CONCURRENCY_PER_CORE = 25; private static final int MIN_CONCURRENCY = 50; private static final int MAX_CONCURRENCY = 1000; private static final long DEFAULT_UPLOAD_PART_COPY_TIMEOUT = TimeUnit.MINUTES.toMillis(2); private final String threadPrefix; + final Logger logger; private final float maxMergeReadSparsityRate; private final int currentIndex; private final Semaphore inflightReadLimiter; @@ -87,6 +102,42 @@ public abstract class AbstractObjectStorage implements ObjectStorage { private final S3LatencyCalculator s3LatencyCalculator; private final Semaphore fastRetryPermit = new Semaphore(MAX_INFLIGHT_FAST_RETRY_COUNT); + /** + * A monitor for successful write requests, in bytes. + */ + private final TrafficMonitor successWriteMonitor = new TrafficMonitor(); + /** + * A monitor for failed write requests (i.e., requests that are throttled), in bytes. + */ + private final TrafficMonitor failedWriteMonitor = new TrafficMonitor(); + /** + * A limiter to control the rate of write requests. + * It is used to limit the write traffic when the write requests are throttled. + */ + private final TrafficRateLimiter writeRateLimiter; + /** + * A limiter to control the volume of write requests. + * It is used to limit the inflight write traffic when the write requests are throttled. + */ + private final TrafficVolumeLimiter writeVolumeLimiter; + /** + * The regulator to control the rate of write requests. + */ + private final TrafficRegulator writeRegulator; + /** + * Pending tasks for write operations (PutObject and UploadPart). + * The task with higher priority and requested earlier will be executed first. + */ + private final Queue writeTasks = new PriorityBlockingQueue<>(); + /** + * The lock to protect the {@link this#currentWriteTask}. + */ + private final Lock writeTaskLock = new ReentrantLock(); + /** + * The current write task (e.g., waiting for {@link this#writeRateLimiter}). + */ + private CompletableFuture currentWriteTask = CompletableFuture.completedFuture(null); + protected AbstractObjectStorage( BucketURI bucketURI, NetworkBandwidthLimiter networkInboundBandwidthLimiter, @@ -98,6 +149,7 @@ protected AbstractObjectStorage( boolean manualMergeRead, String threadPrefix) { this.threadPrefix = threadPrefix; + this.logger = new LogContext(String.format("[ObjectStorage-%s-%s] ", threadPrefix, currentIndex)).logger(AbstractObjectStorage.class); this.bucketURI = bucketURI; this.currentIndex = currentIndex; this.maxMergeReadSparsityRate = Utils.getMaxMergeReadSparsityRate(); @@ -109,18 +161,18 @@ protected AbstractObjectStorage( String prefix = threadPrefix + "-" + currentIndex + "-"; writeLimiterCallbackExecutor = Threads.newFixedThreadPoolWithMonitor(1, - prefix + "s3-write-limiter-cb-executor", true, LOGGER); + prefix + "s3-write-limiter-cb-executor", true, logger); readCallbackExecutor = Threads.newFixedThreadPoolWithMonitor(1, - prefix + "s3-read-cb-executor", true, LOGGER); + prefix + "s3-read-cb-executor", true, logger); writeCallbackExecutor = Threads.newFixedThreadPoolWithMonitor(1, - prefix + "s3-write-cb-executor", true, LOGGER); + prefix + "s3-write-cb-executor", true, logger); scheduler = Threads.newSingleThreadScheduledExecutor( - ThreadUtils.createThreadFactory(prefix + "s3-scheduler", true), LOGGER); + ThreadUtils.createThreadFactory(prefix + "s3-scheduler", true), logger); fastRetryTimer = new HashedWheelTimer( ThreadUtils.createThreadFactory(prefix + "s3-fast-retry-timer", true), 10, TimeUnit.MILLISECONDS, 1000); if (!manualMergeRead) { - scheduler.scheduleWithFixedDelay(this::tryMergeRead, 1, 1, TimeUnit.MILLISECONDS); + scheduler.scheduleWithFixedDelay(this::tryMergeRead, 5, 5, TimeUnit.MILLISECONDS); } S3StreamMetricsManager.registerInflightS3ReadQuotaSupplier(inflightReadLimiter::availablePermits, currentIndex); S3StreamMetricsManager.registerInflightS3WriteQuotaSupplier(inflightWriteLimiter::availablePermits, currentIndex); @@ -133,6 +185,11 @@ protected AbstractObjectStorage( 1024 * 1024, 2 * 1024 * 1024, 3 * 1024 * 1024, 4 * 1024 * 1024, 5 * 1024 * 1024, 8 * 1024 * 1024, 12 * 1024 * 1024, 16 * 1024 * 1024, 32 * 1024 * 1024}, Duration.ofSeconds(3).toMillis()); + + writeRateLimiter = new TrafficRateLimiter(scheduler); + writeVolumeLimiter = new TrafficVolumeLimiter(); + writeRegulator = new TrafficRegulator("write", successWriteMonitor, failedWriteMonitor, writeRateLimiter, writeVolumeLimiter, logger); + scheduler.scheduleWithFixedDelay(writeRegulator::regulate, 60, 60, TimeUnit.SECONDS); } public AbstractObjectStorage(BucketURI bucketURI, @@ -159,7 +216,7 @@ public CompletableFuture rangeRead(ReadOptions options, String objectPa } if (end != RANGE_READ_TO_END && start > end) { IllegalArgumentException ex = new IllegalArgumentException(); - LOGGER.error("[UNEXPECTED] rangeRead [{}, {})", start, end, ex); + logger.error("[UNEXPECTED] rangeRead [{}, {})", start, end, ex); cf.completeExceptionally(ex); return cf; } else if (start == end) { @@ -168,15 +225,7 @@ public CompletableFuture rangeRead(ReadOptions options, String objectPa } BiFunction> networkInboundBandwidthLimiterFunction = - (throttleStrategy, size) -> { - long startTime = System.nanoTime(); - return networkInboundBandwidthLimiter.consume(throttleStrategy, size) - .whenComplete((v, ex) -> - NetworkStats.getInstance() - .networkLimiterQueueTimeStats(AsyncNetworkBandwidthLimiter.Type.INBOUND, throttleStrategy) - .record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS))); - - }; + networkInboundBandwidthLimiter::consume; long acquiredSize = end - start; @@ -203,7 +252,7 @@ public CompletableFuture rangeRead(ReadOptions options, String objectPa }); return FutureUtil.timeoutWithNewReturn(cf, 2, TimeUnit.MINUTES, () -> { - LOGGER.warn("rangeRead {} {}-{} timeout", objectPath, start, end); + logger.warn("rangeRead {} {}-{} timeout", objectPath, start, end); // The return CompletableFuture will be completed with TimeoutException, // so we need to release the ByteBuf if the read complete later. cf.thenAccept(ReferenceCounted::release); @@ -218,18 +267,19 @@ public CompletableFuture write(WriteOptions options, String objectP data.release(); return retCf; } - TimerUtil timerUtil = new TimerUtil(); networkOutboundBandwidthLimiter .consume(options.throttleStrategy(), data.readableBytes()) .whenCompleteAsync((v, ex) -> { - NetworkStats.getInstance().networkLimiterQueueTimeStats(AsyncNetworkBandwidthLimiter.Type.OUTBOUND, options.throttleStrategy()) - .record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); if (ex != null) { - data.release(); cf.completeExceptionally(ex); - } else { - write0(options, objectPath, data, cf); + data.release(); + return; + } + if (checkTimeout(options, cf)) { + data.release(); + return; } + queuedWrite0(options, objectPath, data, cf); }, writeLimiterCallbackExecutor); return retCf; } @@ -238,16 +288,34 @@ private void recordWriteStats(String path, long objectSize, TimerUtil timerUtil) s3LatencyCalculator.record(objectSize, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); S3OperationStats.getInstance().uploadSizeTotalStats.add(MetricsLevel.INFO, objectSize); S3OperationStats.getInstance().putObjectStats(objectSize, true).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("put object {} with size {}, cost {}ms", path, objectSize, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + successWriteMonitor.record(objectSize); + if (logger.isDebugEnabled()) { + logger.debug("put object {} with size {}, cost {}ms", path, objectSize, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); } } - private void write0(WriteOptions options, String path, ByteBuf data, CompletableFuture cf) { + /** + * Put an object to the specified path. + * + * @param options options (or context) about the write operation + * @param path the path to put the object + * @param data the data to put, it will be released once the finalCf is done + * @param attemptCf the CompletableFuture to complete when a single attempt is done + * @param finalCf the CompletableFuture to complete when the write operation is done + */ + private void write0(WriteOptions options, String path, ByteBuf data, CompletableFuture attemptCf, + CompletableFuture finalCf) { TimerUtil timerUtil = new TimerUtil(); long objectSize = data.readableBytes(); + if (checkTimeout(options, finalCf)) { + attemptCf.completeExceptionally(new TimeoutException()); + data.release(); + return; + } + CompletableFuture writeCf = doWrite(options, path, data); + FutureUtil.propagate(writeCf, attemptCf); AtomicBoolean completedFlag = new AtomicBoolean(false); WriteOptions retryOptions = options.copy().retry(true); @@ -267,10 +335,10 @@ private void write0(WriteOptions options, String path, ByteBuf data, Completable fastRetryPermit.release(); if (completedFlag.compareAndSet(false, true)) { - cf.complete(null); - LOGGER.info("Fast retry: put object {} with size {}, cost {}ms, delay {}ms", path, objectSize, retryTimerUtil.elapsedAs(TimeUnit.MILLISECONDS), delayMillis); + finalCf.complete(null); + logger.info("Fast retry: put object {} with size {}, cost {}ms, delay {}ms", path, objectSize, retryTimerUtil.elapsedAs(TimeUnit.MILLISECONDS), delayMillis); } else { - LOGGER.info("Fast retry but duplicated: put object {} with size {}, cost {}ms, delay {}ms", path, objectSize, retryTimerUtil.elapsedAs(TimeUnit.MILLISECONDS), delayMillis); + logger.info("Fast retry but duplicated: put object {} with size {}, cost {}ms, delay {}ms", path, objectSize, retryTimerUtil.elapsedAs(TimeUnit.MILLISECONDS), delayMillis); } }).exceptionally(ignore -> { data.release(); @@ -290,27 +358,57 @@ private void write0(WriteOptions options, String path, ByteBuf data, Completable recordWriteStats(path, objectSize, timerUtil); data.release(); if (completedFlag.compareAndSet(false, true)) { - cf.complete(null); + finalCf.complete(null); } }).exceptionally(ex -> { S3OperationStats.getInstance().putObjectStats(objectSize, false).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); Pair strategyAndCause = toRetryStrategyAndCause(ex, S3Operation.PUT_OBJECT); RetryStrategy retryStrategy = strategyAndCause.getLeft(); Throwable cause = strategyAndCause.getRight(); + if (retryStrategy == RetryStrategy.ABORT || checkS3ApiMode) { - LOGGER.error("PutObject for object {} fail", path, cause); + // no need to retry + logger.error("PutObject for object {} fail", path, cause); data.release(); if (completedFlag.compareAndSet(false, true)) { - cf.completeExceptionally(cause); + finalCf.completeExceptionally(cause); } + return null; + } + + int retryCount = retryOptions.retryCountGetAndAdd(); + if (isThrottled(cause, retryCount)) { + failedWriteMonitor.record(objectSize); + logger.warn("PutObject for object {} fail, retry count {}, queued and retry later", path, retryCount, cause); + queuedWrite0(retryOptions, path, data, finalCf); } else { - LOGGER.warn("PutObject for object {} fail, retry later", path, cause); - scheduler.schedule(() -> write0(retryOptions, path, data, cf), retryDelay(S3Operation.PUT_OBJECT), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.PUT_OBJECT, retryCount); + logger.warn("PutObject for object {} fail, retry count {}, retry in {}ms", path, retryCount, delay, cause); + delayedWrite0(retryOptions, path, data, finalCf, delay); } return null; }); } + private void delayedWrite0(WriteOptions options, String path, ByteBuf data, CompletableFuture cf, + int delayMs) { + CompletableFuture ignored = new CompletableFuture<>(); + scheduler.schedule(() -> write0(options, path, data, ignored, cf), delayMs, TimeUnit.MILLISECONDS); + } + + private void queuedWrite0(WriteOptions options, String path, ByteBuf data, CompletableFuture cf) { + CompletableFuture attemptCf = new CompletableFuture<>(); + AsyncTask task = new AsyncTask( + options.requestTime(), + options.throttleStrategy(), + () -> write0(options, path, data, attemptCf, cf), + attemptCf, + () -> (long) data.readableBytes() + ); + writeTasks.add(task); + maybeRunNextWriteTask(); + } + public CompletableFuture createMultipartUpload(WriteOptions options, String path) { CompletableFuture cf = new CompletableFuture<>(); CompletableFuture retCf = acquireWritePermit(cf); @@ -332,11 +430,12 @@ private void createMultipartUpload0(WriteOptions options, String path, Completab RetryStrategy retryStrategy = strategyAndCause.getLeft(); Throwable cause = strategyAndCause.getRight(); if (retryStrategy == RetryStrategy.ABORT || checkS3ApiMode) { - LOGGER.error("CreateMultipartUpload for object {} fail", path, cause); + logger.error("CreateMultipartUpload for object {} fail", path, cause); cf.completeExceptionally(cause); } else { - LOGGER.warn("CreateMultipartUpload for object {} fail, retry later", path, cause); - scheduler.schedule(() -> createMultipartUpload0(options, path, cf), retryDelay(S3Operation.CREATE_MULTI_PART_UPLOAD), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.CREATE_MULTI_PART_UPLOAD, options.retryCountGetAndAdd()); + logger.warn("CreateMultipartUpload for object {} fail, retry in {}ms", path, delay, cause); + scheduler.schedule(() -> createMultipartUpload0(options, path, cf), delay, TimeUnit.MILLISECONDS); } return null; }); @@ -346,49 +445,99 @@ public CompletableFuture uploadPart(WriteOptions opt int partNumber, ByteBuf data) { CompletableFuture cf = new CompletableFuture<>(); CompletableFuture refCf = acquireWritePermit(cf); + refCf = refCf.whenComplete((v, ex) -> data.release()); if (refCf.isDone()) { - data.release(); return refCf; } networkOutboundBandwidthLimiter .consume(options.throttleStrategy(), data.readableBytes()) .whenCompleteAsync((v, ex) -> { if (ex != null) { - data.release(); cf.completeExceptionally(ex); - } else { - uploadPart0(options, path, uploadId, partNumber, data, cf); + return; } + if (checkTimeout(options, cf)) { + return; + } + queuedUploadPart0(options, path, uploadId, partNumber, data, cf); }, writeLimiterCallbackExecutor); return refCf; } + /** + * Upload a part of an object to the specified path. + * + * @param options options (or context) about the write operation + * @param path the path of the object where the part will be uploaded + * @param uploadId the upload ID of the multipart upload + * @param partNumber the part number of the part to be uploaded + * @param data the data to be uploaded, it will be released once the finalCf is done + * @param attemptCf the CompletableFuture to complete when a single attempt is done + * @param finalCf the CompletableFuture to complete when the upload operation is done + */ private void uploadPart0(WriteOptions options, String path, String uploadId, int partNumber, ByteBuf data, - CompletableFuture cf) { + CompletableFuture attemptCf, + CompletableFuture finalCf) { + if (checkTimeout(options, finalCf)) { + attemptCf.completeExceptionally(new TimeoutException()); + return; + } TimerUtil timerUtil = new TimerUtil(); int size = data.readableBytes(); - doUploadPart(options, path, uploadId, partNumber, data).thenAccept(part -> { + CompletableFuture uploadPartCf = doUploadPart(options, path, uploadId, partNumber, data); + FutureUtil.propagate(uploadPartCf, attemptCf); + uploadPartCf.thenAccept(part -> { S3OperationStats.getInstance().uploadSizeTotalStats.add(MetricsLevel.INFO, size); S3OperationStats.getInstance().uploadPartStats(size, true).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - data.release(); - cf.complete(part); + successWriteMonitor.record(size); + finalCf.complete(part); }).exceptionally(ex -> { S3OperationStats.getInstance().uploadPartStats(size, false).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); Pair strategyAndCause = toRetryStrategyAndCause(ex, S3Operation.UPLOAD_PART); RetryStrategy retryStrategy = strategyAndCause.getLeft(); Throwable cause = strategyAndCause.getRight(); + if (retryStrategy == RetryStrategy.ABORT || checkS3ApiMode) { - LOGGER.error("UploadPart for object {}-{} fail", path, partNumber, cause); - data.release(); - cf.completeExceptionally(cause); + // no need to retry + logger.error("UploadPart for object {}-{} fail", path, partNumber, cause); + finalCf.completeExceptionally(cause); + return null; + } + + int retryCount = options.retryCountGetAndAdd(); + if (isThrottled(cause, retryCount)) { + failedWriteMonitor.record(size); + logger.warn("UploadPart for object {}-{} fail, retry count {}, queued and retry later", path, partNumber, retryCount, cause); + queuedUploadPart0(options, path, uploadId, partNumber, data, finalCf); } else { - LOGGER.warn("UploadPart for object {}-{} fail, retry later", path, partNumber, cause); - scheduler.schedule(() -> uploadPart0(options, path, uploadId, partNumber, data, cf), retryDelay(S3Operation.UPLOAD_PART), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.UPLOAD_PART, retryCount); + logger.warn("UploadPart for object {}-{} fail, retry count {}, retry in {}ms", path, partNumber, retryCount, delay, cause); + delayedUploadPart0(options, path, uploadId, partNumber, data, finalCf, delay); } return null; }); } + private void delayedUploadPart0(WriteOptions options, String path, String uploadId, int partNumber, ByteBuf data, + CompletableFuture cf, int delayMs) { + CompletableFuture ignored = new CompletableFuture<>(); + scheduler.schedule(() -> uploadPart0(options, path, uploadId, partNumber, data, ignored, cf), delayMs, TimeUnit.MILLISECONDS); + } + + private void queuedUploadPart0(WriteOptions options, String path, String uploadId, int partNumber, ByteBuf data, + CompletableFuture cf) { + CompletableFuture attemptCf = new CompletableFuture<>(); + AsyncTask task = new AsyncTask( + options.requestTime(), + options.throttleStrategy(), + () -> uploadPart0(options, path, uploadId, partNumber, data, attemptCf, cf), + attemptCf, + () -> (long) data.readableBytes() + ); + writeTasks.add(task); + maybeRunNextWriteTask(); + } + public CompletableFuture uploadPartCopy(WriteOptions options, String sourcePath, String path, long start, long end, String uploadId, int partNumber) { CompletableFuture cf = new CompletableFuture<>(); @@ -413,13 +562,14 @@ private void uploadPartCopy0(WriteOptions options, String sourcePath, String pat RetryStrategy retryStrategy = strategyAndCause.getLeft(); Throwable cause = strategyAndCause.getRight(); if (retryStrategy == RetryStrategy.ABORT || checkS3ApiMode) { - LOGGER.warn("UploadPartCopy for object {}-{} [{}, {}] fail", path, partNumber, start, end, cause); + logger.warn("UploadPartCopy for object {}-{} [{}, {}] fail", path, partNumber, start, end, cause); cf.completeExceptionally(cause); } else { long nextApiCallAttemptTimeout = Math.min(options.apiCallAttemptTimeout() * 2, TimeUnit.MINUTES.toMillis(10)); - LOGGER.warn("UploadPartCopy for object {}-{} [{}, {}] fail, retry later with apiCallAttemptTimeout={}", path, partNumber, start, end, nextApiCallAttemptTimeout, cause); options.apiCallAttemptTimeout(nextApiCallAttemptTimeout); - scheduler.schedule(() -> uploadPartCopy0(options, sourcePath, path, start, end, uploadId, partNumber, cf), retryDelay(S3Operation.UPLOAD_PART_COPY), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.UPLOAD_PART_COPY, options.retryCountGetAndAdd()); + logger.warn("UploadPartCopy for object {}-{} [{}, {}] fail, retry in {}ms with apiCallAttemptTimeout={}", path, partNumber, start, end, delay, nextApiCallAttemptTimeout, cause); + scheduler.schedule(() -> uploadPartCopy0(options, sourcePath, path, start, end, uploadId, partNumber, cf), delay, TimeUnit.MILLISECONDS); } return null; }); @@ -448,24 +598,26 @@ private void completeMultipartUpload0(WriteOptions options, String path, String RetryStrategy retryStrategy = strategyAndCause.getLeft(); Throwable cause = strategyAndCause.getRight(); if (retryStrategy == RetryStrategy.ABORT || checkS3ApiMode) { - LOGGER.error("CompleteMultipartUpload for object {} fail", path, cause); + logger.error("CompleteMultipartUpload for object {} fail", path, cause); cf.completeExceptionally(cause); } else if (!checkPartNumbers(parts)) { - LOGGER.error("CompleteMultipartUpload for object {} fail, part numbers are not continuous", path); + logger.error("CompleteMultipartUpload for object {} fail, part numbers are not continuous", path); cf.completeExceptionally(new IllegalArgumentException("Part numbers are not continuous")); } else if (retryStrategy == RetryStrategy.VISIBILITY_CHECK) { rangeRead(new ReadOptions().throttleStrategy(ThrottleStrategy.BYPASS).bucket(options.bucketId()), path, 0, 1) .whenComplete((nil, t) -> { if (t != null) { - LOGGER.warn("CompleteMultipartUpload for object {} fail, retry later", path, cause); - scheduler.schedule(() -> completeMultipartUpload0(options, path, uploadId, parts, cf), retryDelay(S3Operation.COMPLETE_MULTI_PART_UPLOAD), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.COMPLETE_MULTI_PART_UPLOAD, options.retryCountGetAndAdd()); + logger.warn("CompleteMultipartUpload for object {} fail, retry in {}ms", path, delay, t); + scheduler.schedule(() -> completeMultipartUpload0(options, path, uploadId, parts, cf), delay, TimeUnit.MILLISECONDS); } else { cf.complete(null); } }); } else { - LOGGER.warn("CompleteMultipartUpload for object {} fail, retry later", path, cause); - scheduler.schedule(() -> completeMultipartUpload0(options, path, uploadId, parts, cf), retryDelay(S3Operation.COMPLETE_MULTI_PART_UPLOAD), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.COMPLETE_MULTI_PART_UPLOAD, options.retryCountGetAndAdd()); + logger.warn("CompleteMultipartUpload for object {} fail, retry in {}ms", path, delay, cause); + scheduler.schedule(() -> completeMultipartUpload0(options, path, uploadId, parts, cf), delay, TimeUnit.MILLISECONDS); } return null; }); @@ -479,7 +631,7 @@ public CompletableFuture delete(List objectPaths) { CompletableFuture cf = new CompletableFuture<>(); for (ObjectPath objectPath : objectPaths) { if (!bucketCheck(objectPath.bucketId(), cf)) { - LOGGER.error("[BUG] {} bucket check fail, expect {}", objectPath, bucketId()); + logger.error("[BUG] {} bucket check fail, expect {}", objectPath, bucketURI.bucketId()); return cf; } } @@ -495,10 +647,10 @@ public CompletableFuture> list(String prefix) { CompletableFuture> cf = doList(prefix); cf.thenAccept(keyList -> { S3OperationStats.getInstance().listObjectsStats(true).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - LOGGER.info("List objects finished, count: {}, cost: {}ms", keyList.size(), timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); + logger.info("List objects finished, count: {}, cost: {}ms", keyList.size(), timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); }).exceptionally(ex -> { S3OperationStats.getInstance().listObjectsStats(false).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); - LOGGER.info("List objects failed, cost: {}, ex: {}", timerUtil.elapsedAs(TimeUnit.NANOSECONDS), ex.getMessage()); + logger.info("List objects failed, cost: {}, ex: {}", timerUtil.elapsedAs(TimeUnit.NANOSECONDS), ex.getMessage()); return null; }); return cf; @@ -506,15 +658,15 @@ public CompletableFuture> list(String prefix) { @Override public short bucketId() { - return bucketURI == null ? 0 : bucketURI.bucketId(); + return bucketURI.bucketId(); } @Override public void close() { - writeLimiterCallbackExecutor.shutdown(); - readCallbackExecutor.shutdown(); - writeCallbackExecutor.shutdown(); - scheduler.shutdown(); + ThreadUtils.shutdownExecutor(writeLimiterCallbackExecutor, 1, TimeUnit.SECONDS); + ThreadUtils.shutdownExecutor(readCallbackExecutor, 1, TimeUnit.SECONDS); + ThreadUtils.shutdownExecutor(writeCallbackExecutor, 1, TimeUnit.SECONDS); + ThreadUtils.shutdownExecutor(scheduler, 1, TimeUnit.SECONDS); fastRetryTimer.stop(); doClose(); } @@ -542,12 +694,12 @@ abstract CompletableFuture doCompleteMultipartUpload(WriteOptions options, abstract CompletableFuture> doList(String prefix); - protected int retryDelay(S3Operation operation) { + protected int retryDelay(S3Operation operation, int retryCount) { switch (operation) { case UPLOAD_PART_COPY: return 1000; default: - return DEFAULT_RETRY_DELAY; + return ThreadLocalRandom.current().nextInt(1000) + Math.min(1000 * (1 << Math.min(retryCount, 16)), (int) (TimeUnit.MINUTES.toMillis(1))); } } @@ -560,12 +712,12 @@ void tryMergeRead() { try { tryMergeRead0(); } catch (Throwable e) { - LOGGER.error("[UNEXPECTED] tryMergeRead fail", e); + logger.error("[UNEXPECTED] tryMergeRead fail", e); } } /** - * Get adjacent read tasks and merge them into one read task which read range is not exceed 16MB. + * Get adjacent read tasks and merge them into one read task which read range is not exceed {@link MergedReadTask#MAX_MERGE_READ_SIZE} (4MB). */ private void tryMergeRead0() { List mergedReadTasks = new ArrayList<>(); @@ -573,6 +725,7 @@ private void tryMergeRead0() { if (waitingReadTasks.isEmpty()) { return; } + waitingReadTasks.sort(Comparator.comparing(ReadTask::objectPath).thenComparingLong(ReadTask::start)); int readPermit = availableReadPermit(); while (readPermit > 0 && !waitingReadTasks.isEmpty()) { Iterator it = waitingReadTasks.iterator(); @@ -600,13 +753,13 @@ private void tryMergeRead0() { mergedReadTasks.forEach( mergedReadTask -> { String path = mergedReadTask.objectPath; - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("merge read: {}, {}-{}, size: {}, sparsityRate: {}", + if (logger.isDebugEnabled()) { + logger.debug("merge read: {}, {}-{}, size: {}, sparsityRate: {}", path, mergedReadTask.start, mergedReadTask.end, mergedReadTask.end - mergedReadTask.start, mergedReadTask.dataSparsityRate); } mergedRangeRead(mergedReadTask.readTasks.get(0).options, path, mergedReadTask.start, mergedReadTask.end) - .whenComplete((rst, ex) -> FutureUtil.suppress(() -> mergedReadTask.handleReadCompleted(rst, ex), LOGGER)); + .whenComplete((rst, ex) -> FutureUtil.suppress(() -> mergedReadTask.handleReadCompleted(rst, ex), logger)); } ); } @@ -632,8 +785,8 @@ private void mergedRangeRead0(ReadOptions options, String path, long start, long doRangeRead(options, path, start, end).thenAccept(buf -> { // the end may be RANGE_READ_TO_END (-1) for read all object long dataSize = buf.readableBytes(); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("GetObject for object {} [{}, {}), size: {}, cost: {} ms", + if (logger.isDebugEnabled()) { + logger.debug("GetObject for object {} [{}, {}), size: {}, cost: {} ms", path, start, end, dataSize, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); } S3OperationStats.getInstance().downloadSizeTotalStats.add(MetricsLevel.INFO, dataSize); @@ -645,23 +798,71 @@ private void mergedRangeRead0(ReadOptions options, String path, long start, long Throwable cause = strategyAndCause.getRight(); if (retryStrategy == RetryStrategy.ABORT || checkS3ApiMode) { if (!(cause instanceof ObjectNotExistException)) { - LOGGER.error("GetObject for object {} [{}, {}) fail", path, start, end, cause); + logger.error("GetObject for object {} [{}, {}) fail", path, start, end, cause); } cf.completeExceptionally(cause); } else { - LOGGER.warn("GetObject for object {} [{}, {}) fail, retry later", path, start, end, cause); - scheduler.schedule(() -> mergedRangeRead0(options, path, start, end, cf), retryDelay(S3Operation.GET_OBJECT), TimeUnit.MILLISECONDS); + int delay = retryDelay(S3Operation.GET_OBJECT, options.retryCountGetAndAdd()); + logger.warn("GetObject for object {} [{}, {}) fail, retry in {}ms", path, start, end, delay, cause); + scheduler.schedule(() -> mergedRangeRead0(options, path, start, end, cf), delay, TimeUnit.MILLISECONDS); } S3OperationStats.getInstance().getObjectStats(size, false).record(timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); return null; }); } + private void maybeRunNextWriteTask() { + writeTaskLock.lock(); + try { + if (!currentWriteTask.isDone()) { + return; + } + + AsyncTask task = writeTasks.poll(); + if (task == null) { + return; + } + + long size = Math.min(task.bytes(), writeRegulator.maxRequestSize()); + currentWriteTask = CompletableFuture.allOf( + writeRateLimiter.consume(size), + writeVolumeLimiter.acquire(size) + ).thenRun(task::run); + task.registerCallback(() -> writeVolumeLimiter.release(size)); + currentWriteTask.whenComplete((nil, ignored) -> maybeRunNextWriteTask()); + } finally { + writeTaskLock.unlock(); + } + } + static int getMaxObjectStorageConcurrency() { int cpuCores = Runtime.getRuntime().availableProcessors(); return Math.max(MIN_CONCURRENCY, Math.min(cpuCores * DEFAULT_CONCURRENCY_PER_CORE, MAX_CONCURRENCY)); } + private static boolean isThrottled(Throwable ex, int retryCount) { + if (ex instanceof S3Exception) { + S3Exception s3Ex = (S3Exception) ex; + return s3Ex.statusCode() == HttpStatusCode.THROTTLING || s3Ex.statusCode() == HttpStatusCode.SERVICE_UNAVAILABLE; + } + // regard timeout as throttled except for the first try + return ex instanceof TimeoutException && retryCount > 0; + } + + /** + * Check whether the operation is timeout, and fail the future with {@link TimeoutException} if timeout. + */ + private static boolean checkTimeout(WriteOptions options, CompletableFuture cf) { + long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - options.requestTime()); + if (elapsedMs > options.timeout()) { + cf.completeExceptionally(new TimeoutException(String.format("request timeout, elapsedMs %d > timeoutMs %d", + elapsedMs, options.timeout()))); + return true; + } else { + return false; + } + } + /** * Acquire read permit, permit will auto release when cf complete. * @@ -811,20 +1012,26 @@ boolean tryMerge(AbstractObjectStorage.ReadTask readTask) { private boolean canMerge(AbstractObjectStorage.ReadTask readTask) { return objectPath != null && - objectPath.equals(readTask.objectPath) && - dataSparsityRate <= this.maxMergeReadSparsityRate && - readTask.end != RANGE_READ_TO_END; + objectPath.equals(readTask.objectPath) && + dataSparsityRate <= this.maxMergeReadSparsityRate && + // Don't allow merge read to end task. + readTask.end != RANGE_READ_TO_END && + end != RANGE_READ_TO_END; } void handleReadCompleted(ByteBuf rst, Throwable ex) { + handleReadCompleted(this.readTasks, this.start, rst, ex); + } + + static void handleReadCompleted(List readTasks, long mergeReadStart, ByteBuf rst, Throwable ex) { if (ex != null) { readTasks.forEach(readTask -> readTask.cf.completeExceptionally(ex)); } else { ArrayList sliceByteBufList = new ArrayList<>(); for (AbstractObjectStorage.ReadTask readTask : readTasks) { - int sliceStart = (int) (readTask.start - start); + int sliceStart = (int) (readTask.start - mergeReadStart); if (readTask.end == RANGE_READ_TO_END) { - sliceByteBufList.add(rst.retainedSlice(sliceStart, rst.readableBytes())); + sliceByteBufList.add(rst.retainedSlice(sliceStart, rst.readableBytes() - sliceStart)); } else { sliceByteBufList.add(rst.retainedSlice(sliceStart, (int) (readTask.end - readTask.start))); } @@ -881,9 +1088,9 @@ public boolean equals(Object obj) { return false; var that = (AbstractObjectStorage.ReadTask) obj; return Objects.equals(this.objectPath, that.objectPath) && - this.start == that.start && - this.end == that.end && - Objects.equals(this.cf, that.cf); + this.start == that.start && + this.end == that.end && + Objects.equals(this.cf, that.cf); } @Override @@ -894,10 +1101,10 @@ public int hashCode() { @Override public String toString() { return "ReadTask[" + - "s3ObjectMetadata=" + objectPath + ", " + - "start=" + start + ", " + - "end=" + end + ", " + - "cf=" + cf + ']'; + "s3ObjectMetadata=" + objectPath + ", " + + "start=" + start + ", " + + "end=" + end + ", " + + "cf=" + cf + ']'; } } @@ -943,4 +1150,60 @@ public String getCheckSum() { return checkSum; } } + + /** + * An object storage operation task. + */ + private static class AsyncTask implements Comparable { + private final long requestTime; + private final ThrottleStrategy strategy; + /** + * A runnable to start the task. + */ + private final Runnable starter; + /** + * A future which will be completed when the task is done (whether success or failure). + */ + private final CompletableFuture finishFuture; + private final Supplier sizeInBytes; + + public AsyncTask(long requestTime, ThrottleStrategy strategy, Runnable starter, + CompletableFuture finishFuture, Supplier sizeInBytes) { + this.requestTime = requestTime; + this.strategy = strategy; + this.starter = starter; + this.finishFuture = finishFuture; + this.sizeInBytes = sizeInBytes; + } + + /** + * Start the async task. + */ + public void run() { + starter.run(); + } + + /** + * Register a callback to be called when the async task is finished. + */ + public void registerCallback(Runnable runnable) { + finishFuture.whenComplete((nil, ignored) -> runnable.run()); + } + + /** + * Get the request size in bytes. + */ + public long bytes() { + return sizeInBytes.get(); + } + + @Override + public int compareTo(AsyncTask other) { + int cmp = this.strategy.compareTo(other.strategy); + if (cmp != 0) { + return cmp; + } + return Long.compare(this.requestTime, other.requestTime); + } + } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/AutoMQStaticCredentialsProvider.java b/s3stream/src/main/java/com/automq/stream/s3/operator/AutoMQStaticCredentialsProvider.java new file mode 100644 index 0000000000..e0d1cacc39 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/AutoMQStaticCredentialsProvider.java @@ -0,0 +1,34 @@ +package com.automq.stream.s3.operator; + +import org.apache.commons.lang3.StringUtils; + +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; + +public class AutoMQStaticCredentialsProvider implements AwsCredentialsProvider { + // Deprecated. Use AWS_ACCESS_KEY_ID instead. + private static final String KAFKA_S3_ACCESS_KEY = "KAFKA_S3_ACCESS_KEY"; + // Deprecated. Use AWS_SECRET_ACCESS_KEY instead. + private static final String KAFKA_S3_SECRET_KEY = "KAFKA_S3_SECRET_KEY"; + private final AwsCredentialsProvider staticCredentialsProvider; + + public AutoMQStaticCredentialsProvider(BucketURI bucketURI) { + String accessKey = bucketURI.extensionString(BucketURI.ACCESS_KEY_KEY, System.getenv(KAFKA_S3_ACCESS_KEY)); + String secretKey = bucketURI.extensionString(BucketURI.SECRET_KEY_KEY, System.getenv(KAFKA_S3_SECRET_KEY)); + if (StringUtils.isBlank(accessKey) || StringUtils.isBlank(secretKey)) { + staticCredentialsProvider = null; + return; + } + staticCredentialsProvider = StaticCredentialsProvider.create(AwsBasicCredentials.create(accessKey, secretKey)); + } + + @Override + public AwsCredentials resolveCredentials() { + if (staticCredentialsProvider == null) { + throw new RuntimeException("AK/SK not set in bucket URI"); + } + return staticCredentialsProvider.resolveCredentials(); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/AwsObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/AwsObjectStorage.java index fa4aa9b7a3..7c42c13d43 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/AwsObjectStorage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/AwsObjectStorage.java @@ -1,17 +1,24 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; -import com.automq.stream.s3.ByteBufAlloc; import com.automq.stream.s3.exceptions.ObjectNotExistException; import com.automq.stream.s3.metrics.operations.S3Operation; import com.automq.stream.s3.network.NetworkBandwidthLimiter; @@ -19,18 +26,20 @@ import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.net.URI; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; import java.util.Date; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeoutException; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -39,15 +48,14 @@ import io.netty.buffer.Unpooled; import io.netty.handler.ssl.OpenSsl; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; -import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.auth.credentials.AwsCredentialsProviderChain; -import software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider; -import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.async.AsyncResponseTransformer; import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.core.exception.ApiCallAttemptTimeoutException; import software.amazon.awssdk.core.exception.SdkClientException; import software.amazon.awssdk.http.HttpStatusCode; import software.amazon.awssdk.http.async.SdkAsyncHttpClient; @@ -84,11 +92,10 @@ @SuppressWarnings({"this-escape", "NPathComplexity"}) public class AwsObjectStorage extends AbstractObjectStorage { + // use the root logger to log the error to both log file and stdout + private static final Logger READINESS_CHECK_LOGGER = LoggerFactory.getLogger("ObjectStorageReadinessCheck"); public static final String S3_API_NO_SUCH_KEY = "NoSuchKey"; public static final String PATH_STYLE_KEY = "pathStyle"; - public static final String AUTH_TYPE_KEY = "authType"; - public static final String STATIC_AUTH_TYPE = "static"; - public static final String INSTANCE_AUTH_TYPE = "instance"; public static final String CHECKSUM_ALGORITHM_KEY = "checksumAlgorithm"; // https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html @@ -102,8 +109,6 @@ public class AwsObjectStorage extends AbstractObjectStorage { private final ChecksumAlgorithm checksumAlgorithm; - private static volatile InstanceProfileCredentialsProvider instanceProfileCredentialsProvider; - public AwsObjectStorage(BucketURI bucketURI, Map tagging, NetworkBandwidthLimiter networkInboundBandwidthLimiter, NetworkBandwidthLimiter networkOutboundBandwidthLimiter, boolean readWriteIsolate, boolean checkMode, String threadPrefix) { @@ -123,7 +128,10 @@ public AwsObjectStorage(BucketURI bucketURI, Map tagging, } this.checksumAlgorithm = checksumAlgorithm; - Supplier clientSupplier = () -> newS3Client(bucketURI.endpoint(), bucketURI.region(), bucketURI.extensionBool(PATH_STYLE_KEY, false), credentialsProviders, getMaxObjectStorageConcurrency()); + long apiCallTimeoutMs = Long.parseLong(bucketURI.extensionString(BucketURI.API_CALL_TIMEOUT_KEY, "30000")); + long apiCallAttemptTimeoutMs = Long.parseLong(bucketURI.extensionString(BucketURI.API_CALL_ATTEMPT_TIMEOUT_KEY, "10000")); + + Supplier clientSupplier = () -> newS3Client(bucketURI.endpoint(), bucketURI.region(), bucketURI.extensionBool(PATH_STYLE_KEY, false), credentialsProviders, getMaxObjectStorageConcurrency(), apiCallTimeoutMs, apiCallAttemptTimeoutMs); this.writeS3Client = clientSupplier.get(); this.readS3Client = readWriteIsolate ? clientSupplier.get() : writeS3Client; } @@ -142,7 +150,7 @@ public static Builder builder() { return new Builder(); } - static void checkDeleteObjectsResponse(DeleteObjectsResponse response) throws Exception { + void checkDeleteObjectsResponse(DeleteObjectsResponse response) throws Exception { int errDeleteCount = 0; ArrayList failedKeys = new ArrayList<>(); ArrayList errorsMessages = new ArrayList<>(); @@ -152,7 +160,7 @@ static void checkDeleteObjectsResponse(DeleteObjectsResponse response) throws Ex continue; } if (errDeleteCount < 5) { - LOGGER.error("Delete objects for key [{}] error code [{}] message [{}]", + logger.error("Delete objects for key [{}] error code [{}] message [{}]", error.key(), error.code(), error.message()); } failedKeys.add(error.key()); @@ -175,7 +183,8 @@ CompletableFuture doRangeRead(ReadOptions options, String path, long st CompletableFuture cf = new CompletableFuture<>(); readS3Client.getObject(builder.build(), AsyncResponseTransformer.toPublisher()) .thenAccept(responsePublisher -> { - CompositeByteBuf buf = ByteBufAlloc.compositeByteBuffer(); + // Set maxNumComponents to Integer.MAX_VALUE to avoid #consolidateIfNeeded causing a GC issue. + CompositeByteBuf buf = Unpooled.compositeBuffer(Integer.MAX_VALUE); responsePublisher.subscribe(bytes -> { // the aws client will copy DefaultHttpContent to heap ByteBuffer buf.addComponent(true, Unpooled.wrappedBuffer(bytes)); @@ -324,7 +333,6 @@ Pair toRetryStrategyAndCause(Throwable ex, S3Operation if (cause instanceof S3Exception) { S3Exception s3Ex = (S3Exception) cause; switch (s3Ex.statusCode()) { - case HttpStatusCode.FORBIDDEN: case HttpStatusCode.NOT_FOUND: strategy = RetryStrategy.ABORT; break; @@ -341,6 +349,8 @@ Pair toRetryStrategyAndCause(Throwable ex, S3Operation cause = new ObjectNotExistException(cause); } } + } else if (cause instanceof ApiCallAttemptTimeoutException) { + cause = new TimeoutException(cause.getMessage()); } return Pair.of(strategy, cause); } @@ -355,12 +365,33 @@ void doClose() { @Override CompletableFuture> doList(String prefix) { - return readS3Client.listObjectsV2(builder -> builder.bucket(bucket).prefix(prefix)) - .thenApply(resp -> - resp.contents() - .stream() - .map(object -> new ObjectInfo(bucketURI.bucketId(), object.key(), object.lastModified().toEpochMilli(), object.size())) - .collect(Collectors.toList())); + CompletableFuture> resultFuture = new CompletableFuture<>(); + List allObjects = new ArrayList<>(); + listNextBatch(prefix, null, allObjects, resultFuture); + return resultFuture; + } + + private void listNextBatch(String prefix, String continuationToken, List allObjects, + CompletableFuture> resultFuture) { + readS3Client.listObjectsV2(builder -> { + builder.bucket(bucket).prefix(prefix); + if (continuationToken != null) { + builder.continuationToken(continuationToken); + } + }).thenAccept(resp -> { + resp.contents() + .stream() + .map(object -> new ObjectInfo(bucketURI.bucketId(), object.key(), object.lastModified().toEpochMilli(), object.size())) + .forEach(allObjects::add); + if (resp.isTruncated()) { + listNextBatch(prefix, resp.nextContinuationToken(), allObjects, resultFuture); + } else { + resultFuture.complete(allObjects); + } + }).exceptionally(ex -> { + resultFuture.completeExceptionally(ex); + return null; + }); } @Override @@ -369,33 +400,11 @@ protected DeleteObjectsAccumulator newDeleteObjectsAccumulator() { } protected List credentialsProviders() { - String authType = bucketURI.extensionString(AUTH_TYPE_KEY, STATIC_AUTH_TYPE); - switch (authType) { - case STATIC_AUTH_TYPE: { - String accessKey = bucketURI.extensionString(BucketURI.ACCESS_KEY_KEY, System.getenv("KAFKA_S3_ACCESS_KEY")); - String secretKey = bucketURI.extensionString(BucketURI.SECRET_KEY_KEY, System.getenv("KAFKA_S3_SECRET_KEY")); - if (StringUtils.isBlank(accessKey) || StringUtils.isBlank(secretKey)) { - return Collections.emptyList(); - } - return List.of(StaticCredentialsProvider.create(AwsBasicCredentials.create(accessKey, secretKey))); - } - case INSTANCE_AUTH_TYPE: { - return List.of(instanceProfileCredentialsProvider()); - } - default: - throw new UnsupportedOperationException("Unsupported auth type: " + authType); - } + return credentialsProviders0(bucketURI); } - protected AwsCredentialsProvider instanceProfileCredentialsProvider() { - if (instanceProfileCredentialsProvider == null) { - synchronized (AwsObjectStorage.class) { - if (instanceProfileCredentialsProvider == null) { - instanceProfileCredentialsProvider = InstanceProfileCredentialsProvider.builder().build(); - } - } - } - return instanceProfileCredentialsProvider; + protected List credentialsProviders0(BucketURI bucketURI) { + return List.of(new AutoMQStaticCredentialsProvider(bucketURI), DefaultCredentialsProvider.create()); } private String range(long start, long end) { @@ -407,13 +416,13 @@ private String range(long start, long end) { } protected S3AsyncClient newS3Client(String endpoint, String region, boolean forcePathStyle, - List credentialsProviders, int maxConcurrency) { + List credentialsProviders, int maxConcurrency, long apiCallTimeoutMs, long apiCallAttemptTimeoutMs) { S3AsyncClientBuilder builder = S3AsyncClient.builder().region(Region.of(region)); if (StringUtils.isNotBlank(endpoint)) { builder.endpointOverride(URI.create(endpoint)); } if (!OpenSsl.isAvailable()) { - LOGGER.warn("OpenSSL is not available, using JDK SSL provider, which may have performance issue.", OpenSsl.unavailabilityCause()); + logger.warn("OpenSSL is not available, using JDK SSL provider, which may have performance issue.", OpenSsl.unavailabilityCause()); } SdkAsyncHttpClient httpClient = NettyNioAsyncHttpClient.builder() .maxConcurrency(maxConcurrency) @@ -421,21 +430,20 @@ protected S3AsyncClient newS3Client(String endpoint, String region, boolean forc builder.httpClient(httpClient); builder.serviceConfiguration(c -> c.pathStyleAccessEnabled(forcePathStyle)); builder.credentialsProvider(newCredentialsProviderChain(credentialsProviders)); - builder.overrideConfiguration(clientOverrideConfiguration()); + builder.overrideConfiguration(clientOverrideConfiguration(apiCallTimeoutMs, apiCallAttemptTimeoutMs)); return builder.build(); } - protected ClientOverrideConfiguration clientOverrideConfiguration() { + protected ClientOverrideConfiguration clientOverrideConfiguration(long apiCallTimeoutMs, long apiCallAttemptTimeoutMs) { return ClientOverrideConfiguration.builder() - .apiCallTimeout(Duration.ofMinutes(2)) - .apiCallAttemptTimeout(Duration.ofSeconds(60)) + .apiCallTimeout(Duration.ofMillis(apiCallTimeoutMs)) + .apiCallAttemptTimeout(Duration.ofMillis(apiCallAttemptTimeoutMs)) .build(); } - private AwsCredentialsProvider newCredentialsProviderChain(List credentialsProviders) { + protected AwsCredentialsProvider newCredentialsProviderChain(List credentialsProviders) { List providers = new ArrayList<>(credentialsProviders); // Add default providers to the end of the chain - providers.add(InstanceProfileCredentialsProvider.create()); providers.add(AnonymousCredentialsProvider.create()); return AwsCredentialsProviderChain.builder() .reuseLastProviderEnabled(true) @@ -449,25 +457,24 @@ public boolean readinessCheck() { class ReadinessCheck { public boolean readinessCheck() { - LOGGER.info("Start readiness check for {}", bucketURI); + READINESS_CHECK_LOGGER.info("Start readiness check for {}", bucketURI); String normalPath = String.format("__automq/readiness_check/normal_obj/%d", System.nanoTime()); try { writeS3Client.headObject(HeadObjectRequest.builder().bucket(bucket).key(normalPath).build()).get(); } catch (Throwable e) { - // 权限 / endpoint / xxx Throwable cause = FutureUtil.cause(e); if (cause instanceof SdkClientException) { - LOGGER.error("Cannot connect to s3, please check the s3 endpoint config", cause); + READINESS_CHECK_LOGGER.error("Cannot connect to s3, please check the s3 endpoint config", cause); } else if (cause instanceof S3Exception) { int code = ((S3Exception) cause).statusCode(); switch (code) { case HttpStatusCode.NOT_FOUND: break; case HttpStatusCode.FORBIDDEN: - LOGGER.error("Please check whether config is correct", cause); + READINESS_CHECK_LOGGER.error("Please check whether config is correct", cause); return false; default: - LOGGER.error("Please check config is correct", cause); + READINESS_CHECK_LOGGER.error("Please check config is correct", cause); } } } @@ -478,9 +485,9 @@ public boolean readinessCheck() { } catch (Throwable e) { Throwable cause = FutureUtil.cause(e); if (cause instanceof S3Exception && ((S3Exception) cause).statusCode() == HttpStatusCode.NOT_FOUND) { - LOGGER.error("Cannot find the bucket={}", bucket, cause); + READINESS_CHECK_LOGGER.error("Cannot find the bucket={}", bucket, cause); } else { - LOGGER.error("Please check the identity have the permission to do Write Object operation", cause); + READINESS_CHECK_LOGGER.error("Please check the identity have the permission to do Write Object operation", cause); } return false; } @@ -488,7 +495,7 @@ public boolean readinessCheck() { try { doDeleteObjects(List.of(normalPath)).get(); } catch (Throwable e) { - LOGGER.error("Please check the identity have the permission to do Delete Object operation", FutureUtil.cause(e)); + READINESS_CHECK_LOGGER.error("Please check the identity have the permission to do Delete Object operation", FutureUtil.cause(e)); return false; } @@ -505,15 +512,15 @@ public boolean readinessCheck() { buf.readBytes(readContent); buf.release(); if (!Arrays.equals(content, readContent)) { - LOGGER.error("Read get mismatch content from multi-part upload object, expect {}, but {}", content, readContent); + READINESS_CHECK_LOGGER.error("Read get mismatch content from multi-part upload object, expect {}, but {}", content, readContent); } doDeleteObjects(List.of(multiPartPath)).get(); } catch (Throwable e) { - LOGGER.error("Please check the identity have the permission to do MultiPart Object operation", FutureUtil.cause(e)); + READINESS_CHECK_LOGGER.error("Please check the identity have the permission to do MultiPart Object operation", FutureUtil.cause(e)); return false; } - LOGGER.info("Readiness check pass!"); + READINESS_CHECK_LOGGER.info("Readiness check pass!"); return true; } diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/BucketURI.java b/s3stream/src/main/java/com/automq/stream/s3/operator/BucketURI.java index cd5a3f631b..f101339273 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/BucketURI.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/BucketURI.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -32,6 +40,8 @@ public class BucketURI { private static final String REGION_KEY = "region"; public static final String ACCESS_KEY_KEY = "accessKey"; public static final String SECRET_KEY_KEY = "secretKey"; + public static final String API_CALL_TIMEOUT_KEY = "apiCallTimeoutMs"; + public static final String API_CALL_ATTEMPT_TIMEOUT_KEY = "apiCallAttemptTimeoutMs"; private static final String EMPTY_STRING = ""; private final short bucketId; private final String protocol; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/DeleteObjectsAccumulator.java b/s3stream/src/main/java/com/automq/stream/s3/operator/DeleteObjectsAccumulator.java index 7aa9bcfbdb..172982ff35 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/DeleteObjectsAccumulator.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/DeleteObjectsAccumulator.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/LocalFileObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/LocalFileObjectStorage.java new file mode 100644 index 0000000000..dacb3c45ff --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/LocalFileObjectStorage.java @@ -0,0 +1,418 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import com.automq.stream.s3.exceptions.ObjectNotExistException; +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Threads; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.DirectoryNotEmptyException; +import java.nio.file.DirectoryStream; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.NoSuchFileException; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.StandardCopyOption; +import java.nio.file.StandardOpenOption; +import java.nio.file.attribute.BasicFileAttributes; +import java.util.ArrayList; +import java.util.LinkedList; +import java.util.List; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicLong; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class LocalFileObjectStorage implements ObjectStorage { + public static final short BUCKET_ID = -2; + private static final Logger LOGGER = LoggerFactory.getLogger(LocalFileObjectStorage.class); + private final BucketURI bucketURI; + private final Path dataParentPath; + private final String dataParentPathStr; + private final String atomicWriteParentPathStr; + private final AtomicLong atomicWriteCounter = new AtomicLong(0); + // thread-safe is guarded by synchronized(availableSpace) + final AtomicLong availableSpace = new AtomicLong(); + final Queue waitingTasks = new LinkedBlockingQueue<>(); + private final ExecutorService ioExecutor = Threads.newFixedThreadPoolWithMonitor(8, "LOCAL_FILE_OBJECT_STORAGE_IO", true, LOGGER); + + public LocalFileObjectStorage(BucketURI bucketURI) { + if (bucketURI.bucketId() != BUCKET_ID) { + throw new IllegalArgumentException("bucketId must be -2"); + } + this.bucketURI = bucketURI; + this.dataParentPathStr = bucketURI.bucket() + File.separator + "data"; + this.dataParentPath = Path.of(dataParentPathStr); + this.atomicWriteParentPathStr = bucketURI.bucket() + File.separator + "atomic"; + + try { + if (!dataParentPath.toFile().isDirectory()) { + Files.createDirectories(dataParentPath); + } + availableSpace.set(new File(dataParentPathStr).getFreeSpace() - 2L * 1024 * 1024 * 1024); + } catch (IOException e) { + throw new RuntimeException(e); + } + try (DirectoryStream stream = Files.newDirectoryStream(Path.of(atomicWriteParentPathStr))) { + for (Path file : stream) { + if (Files.isRegularFile(file)) { + Files.delete(file); + } + } + } catch (IOException ignored) { + } + } + + @Override + public boolean readinessCheck() { + return true; + } + + @Override + public void close() { + + } + + @Override + public Writer writer(WriteOptions options, String objectPath) { + options.bucketId(bucketId()); + return new LocalFileWriter(objectPath); + } + + @Override + public CompletableFuture rangeRead(ReadOptions options, String objectPath, long start, long end) { + CompletableFuture cf = new CompletableFuture<>(); + Path path = dataPath(objectPath); + ioExecutor.submit(() -> { + long position = start; + try (FileChannel fileChannel = FileChannel.open(path, StandardOpenOption.READ)) { + byte[] bytes; + if (end == -1) { + bytes = new byte[(int) (fileChannel.size() - start)]; + } else { + bytes = new byte[(int) (end - start)]; + } + ByteBuffer buffer = ByteBuffer.wrap(bytes); + while (buffer.hasRemaining()) { + int readSize = fileChannel.read(buffer, position); + if (readSize == -1) { + cf.completeExceptionally(new IllegalArgumentException(String.format("rangeRead %s [%s, %s) out of bound [0, %s)", + objectPath, start, end, fileChannel.size()))); + return; + } + position += readSize; + } + cf.complete(Unpooled.wrappedBuffer(bytes)); + } catch (NoSuchFileException e) { + cf.completeExceptionally(new ObjectNotExistException()); + } catch (Throwable e) { + cf.completeExceptionally(e); + } + }); + return cf; + } + + @Override + public CompletableFuture> list(String prefix) { + CompletableFuture> cf = new CompletableFuture<>(); + try { + Path path = dataPath(prefix); + String pathPrefix = path.toString(); + if (!Files.isDirectory(path)) { + path = path.getParent(); + } + List list = new ArrayList<>(); + if (path != null && Files.exists(path)) { + Files.walkFileTree(path, new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) { + if (!Files.isDirectory(file)) { + String filePathStr = file.toString(); + if (filePathStr.startsWith(pathPrefix)) { + list.add(new ObjectInfo(bucketId(), filePathStr.substring(dataParentPathStr.length() + 1), attrs.creationTime().toMillis(), attrs.size())); + } + } + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) { + return FileVisitResult.CONTINUE; + } + }); + } + cf.complete(list); + } catch (Throwable e) { + cf.completeExceptionally(e); + } + return cf; + } + + @Override + public CompletableFuture delete(List objectPaths) { + CompletableFuture cf = new CompletableFuture<>(); + ioExecutor.submit(() -> { + long size = 0; + try { + for (ObjectPath objectPath : objectPaths) { + Path dataPath = dataPath(objectPath.key()); + long fileSize = fileSize(dataPath); + deleteFileAndEmptyParents(dataPath); + size += fileSize; + } + cf.complete(null); + } catch (Throwable e) { + cf.completeExceptionally(e); + } finally { + freeSpace(size); + } + }); + return cf; + } + + @Override + public short bucketId() { + return bucketURI.bucketId(); + } + + private Path atomicWritePath() { + return Path.of(atomicWriteParentPathStr + File.separator + atomicWriteCounter.incrementAndGet()); + } + + private Path dataPath(String objectPath) { + return Path.of(dataParentPathStr + File.separator + objectPath); + } + + public void deleteFileAndEmptyParents(Path filePath) throws IOException { + if (!Files.deleteIfExists(filePath)) { + return; + } + + Path parentDir = filePath.getParent(); + + while (parentDir != null && !parentDir.equals(dataParentPath)) { + try { + Files.delete(parentDir); + parentDir = parentDir.getParent(); + } catch (DirectoryNotEmptyException | NoSuchFileException e) { + break; + } + } + } + + private void acquireSpace(int size, Runnable task) { + boolean acquired = false; + synchronized (availableSpace) { + if (availableSpace.get() > size) { + availableSpace.addAndGet(-size); + acquired = true; + } else { + waitingTasks.add(new WriteTask(size, task)); + LOGGER.info("[LOCAL_FILE_OBJECT_STORAGE_FULL]"); + } + } + if (acquired) { + task.run(); + } + } + + private void freeSpace(long size) { + synchronized (availableSpace) { + availableSpace.addAndGet(size); + for (;;) { + WriteTask task = waitingTasks.peek(); + if (task == null) { + break; + } + if (task.writeSize > availableSpace.get()) { + break; + } + waitingTasks.poll(); + availableSpace.addAndGet(-task.writeSize); + FutureUtil.suppress(task.task::run, LOGGER); + } + } + } + + class LocalFileWriter implements Writer { + private final Path atomicWritePath; + private final Path dataPath; + private final FileChannel fileChannel; + private Throwable cause; + private final List> writeCfList = new LinkedList<>(); + private long nextWritePosition = 0; + + public LocalFileWriter(String objectPath) { + this.atomicWritePath = atomicWritePath(); + this.dataPath = dataPath(objectPath); + FileChannel fileChannel = null; + try { + for (int i = 0; i < 2; i++) { + try { + fileChannel = FileChannel.open(atomicWritePath, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING, StandardOpenOption.WRITE); + } catch (NoSuchFileException e) { + Path parent = atomicWritePath.getParent(); + if (parent != null) { + Files.createDirectories(parent); + } + } + } + if (fileChannel == null) { + throw new IllegalStateException("expect file channel create success"); + } + } catch (Throwable e) { + cause = e; + } + this.fileChannel = fileChannel; + } + + @Override + public CompletableFuture write(ByteBuf data) { + CompletableFuture cf = new CompletableFuture<>(); + CompletableFuture retCf = cf.whenComplete((v, ex) -> data.release()); + if (cause != null) { + cf.completeExceptionally(cause); + return retCf; + } + long startWritePosition = nextWritePosition; + int dataSize = data.readableBytes(); + nextWritePosition += dataSize; + acquireSpace(dataSize, () -> ioExecutor.execute(() -> { + long position = startWritePosition; + try { + ByteBuffer[] buffers = data.nioBuffers(); + for (ByteBuffer buf : buffers) { + while (buf.hasRemaining()) { + position += fileChannel.write(buf, position); + } + } + cf.complete(null); + } catch (Throwable e) { + cf.completeExceptionally(e); + } + })); + writeCfList.add(retCf); + return retCf; + } + + @Override + public void copyOnWrite() { + throw new UnsupportedOperationException(); + } + + @Override + public void copyWrite(S3ObjectMetadata s3ObjectMetadata, long start, long end) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean hasBatchingPart() { + return false; + } + + @Override + public CompletableFuture close() { + CompletableFuture cf = new CompletableFuture<>(); + CompletableFuture.allOf(writeCfList.toArray(new CompletableFuture[0])).whenComplete((v, ex) -> { + if (ex != null) { + cf.completeExceptionally(ex); + return; + } + ioExecutor.execute(() -> { + try (fileChannel) { + fileChannel.force(true); + try { + Files.move(atomicWritePath, dataPath, StandardCopyOption.REPLACE_EXISTING); + } catch (NoSuchFileException e) { + Path parent = dataPath.getParent(); + if (parent != null) { + Files.createDirectories(parent); + } + Files.move(atomicWritePath, dataPath, StandardCopyOption.REPLACE_EXISTING); + } + cf.complete(null); + } catch (Throwable e) { + cf.completeExceptionally(e); + } + }); + }); + return cf; + } + + @Override + public CompletableFuture release() { + return CompletableFuture.allOf(writeCfList.toArray(new CompletableFuture[0])); + } + + @Override + public short bucketId() { + return LocalFileObjectStorage.this.bucketId(); + } + } + + public static Builder builder() { + return new Builder(); + } + + private static long fileSize(Path filePath) { + try { + return Files.size(filePath); + } catch (IOException e) { + return 0; + } + } + + public static class Builder { + private BucketURI bucketURI; + + public Builder bucket(BucketURI bucketURI) { + this.bucketURI = bucketURI; + return this; + } + + public LocalFileObjectStorage build() { + return new LocalFileObjectStorage(bucketURI); + } + } + + static class WriteTask { + final long writeSize; + // The task should be a none blocking task + final Runnable task; + + public WriteTask(long writeSize, Runnable task) { + this.writeSize = writeSize; + this.task = task; + } + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/MemoryObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/MemoryObjectStorage.java index 8a6b66d7e6..b51afafd64 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/MemoryObjectStorage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/MemoryObjectStorage.java @@ -1,17 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.exceptions.ObjectNotExistException; import com.automq.stream.s3.metadata.S3ObjectMetadata; import com.automq.stream.s3.metrics.operations.S3Operation; import com.automq.stream.s3.network.NetworkBandwidthLimiter; @@ -47,6 +56,13 @@ public MemoryObjectStorage(boolean manualMergeRead, short bucketId) { this.bucketId = bucketId; } + public MemoryObjectStorage(int concurrencyCount) { + super(BucketURI.parse(0 + "@s3://b"), + new RecordTestNetworkBandwidthLimiter(), new RecordTestNetworkBandwidthLimiter(), + concurrencyCount, 0, true, false, false, "memory"); + this.bucketId = 0; + } + public MemoryObjectStorage(short bucketId) { this(false, bucketId); } @@ -63,7 +79,7 @@ public MemoryObjectStorage(boolean manualMergeRead) { CompletableFuture doRangeRead(ReadOptions options, String path, long start, long end) { ByteBuf value = storage.get(path); if (value == null) { - return FutureUtil.failedFuture(new IllegalArgumentException("object not exist")); + return FutureUtil.failedFuture(new ObjectNotExistException("object not exist")); } int length = end != -1L ? (int) (end - start) : (int) (value.readableBytes() - start); ByteBuf rst = value.retainedSlice(value.readerIndex() + (int) start, length); @@ -185,7 +201,7 @@ public Set getDeleteObjectKeys() { @Override Pair toRetryStrategyAndCause(Throwable ex, S3Operation operation) { Throwable cause = FutureUtil.cause(ex); - RetryStrategy strategy = cause instanceof UnsupportedOperationException || cause instanceof IllegalArgumentException + RetryStrategy strategy = (cause instanceof UnsupportedOperationException || cause instanceof IllegalArgumentException || cause instanceof ObjectNotExistException) ? RetryStrategy.ABORT : RetryStrategy.RETRY; return Pair.of(strategy, cause); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/MultiPartWriter.java b/s3stream/src/main/java/com/automq/stream/s3/operator/MultiPartWriter.java index f8f535798b..14cbde2c78 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/MultiPartWriter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/MultiPartWriter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -32,7 +40,7 @@ import io.netty.buffer.CompositeByteBuf; public class MultiPartWriter implements Writer { - private static final long MAX_MERGE_WRITE_SIZE = 16L * 1024 * 1024; + private static final long MAX_MERGE_WRITE_SIZE = 32L * 1024 * 1024; final CompletableFuture uploadIdCf = new CompletableFuture<>(); private final AbstractObjectStorage operator; private final String path; @@ -43,6 +51,7 @@ public class MultiPartWriter implements Writer { * The minPartSize represents the minimum size of a part for a multipart object. */ private final long minPartSize; + private final long maxMergeWriteSize; private final TimerUtil timerUtil = new TimerUtil(); private final AtomicLong totalWriteSize = new AtomicLong(0L); private String uploadId; @@ -51,10 +60,16 @@ public class MultiPartWriter implements Writer { public MultiPartWriter(ObjectStorage.WriteOptions writeOptions, AbstractObjectStorage operator, String path, long minPartSize) { + this(writeOptions, operator, path, minPartSize, MAX_MERGE_WRITE_SIZE); + } + + public MultiPartWriter(ObjectStorage.WriteOptions writeOptions, AbstractObjectStorage operator, String path, + long minPartSize, long maxMergeWriteSize) { this.writeOptions = writeOptions; this.operator = operator; this.path = path; this.minPartSize = minPartSize; + this.maxMergeWriteSize = maxMergeWriteSize; init(); } @@ -78,7 +93,7 @@ public CompletableFuture write(ByteBuf data) { ObjectPart objectPart = this.objectPart; objectPart.write(data); - if (objectPart.size() > minPartSize) { + if (objectPart.size() > Math.max(minPartSize, maxMergeWriteSize)) { objectPart.upload(); // finish current part. this.objectPart = null; @@ -121,7 +136,7 @@ public void copyWrite0(S3ObjectMetadata sourceObjectMateData, long start, long e new CopyObjectPart(sourceObjectMateData.key(), start, end); } } else { - if (objectPart.size() + targetSize > MAX_MERGE_WRITE_SIZE) { + if (objectPart.size() + targetSize > maxMergeWriteSize) { long readAndWriteCopyEnd = start + minPartSize - objectPart.size(); objectPart.readAndWrite(sourceObjectMateData, start, readAndWriteCopyEnd); objectPart.upload(); @@ -163,8 +178,13 @@ public CompletableFuture close() { @Override public CompletableFuture release() { + List> partsToWait = parts; + if (objectPart != null) { + // skip waiting for pending part + partsToWait = partsToWait.subList(0, partsToWait.size() - 1); + } // wait for all ongoing uploading parts to finish and release pending part - return CompletableFuture.allOf(parts.toArray(new CompletableFuture[0])).whenComplete((nil, ex) -> { + return CompletableFuture.allOf(partsToWait.toArray(new CompletableFuture[0])).whenComplete((nil, ex) -> { if (objectPart != null) { objectPart.release(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java index 4463d047d6..bbc47f3e02 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -46,18 +54,24 @@ default CompletableFuture read(ReadOptions options, String objectPath) /** * Range read object from the object storage. - * It will throw {@link ObjectNotExistException} if the object not found. + * It will failFuture with {@link ObjectNotExistException} if the object not found. + * @param options {@link ReadOptions} + * @param objectPath the object path + * @param start inclusive start position + * @param end exclusive end position + * @return read result */ CompletableFuture rangeRead(ReadOptions options, String objectPath, long start, long end); // Low level API - CompletableFuture write(WriteOptions options, String objectPath, ByteBuf buf); + default CompletableFuture write(WriteOptions options, String objectPath, ByteBuf buf) { + Writer writer = writer(options, objectPath); + writer.write(buf); + return writer.close().thenApply(nil -> new WriteResult(bucketId())); + } CompletableFuture> list(String prefix); - // NOTE: this is a temporary method to get bucketId for direct read with object storage interface - short bucketId(); - /** * The deleteObjects API have max batch limit. * see @@ -67,6 +81,8 @@ default CompletableFuture read(ReadOptions options, String objectPath) */ CompletableFuture delete(List objectPaths); + short bucketId(); + class ObjectPath { private final short bucketId; private final String key; @@ -117,10 +133,17 @@ class WriteOptions { private ThrottleStrategy throttleStrategy = ThrottleStrategy.BYPASS; private int allocType = ByteBufAlloc.DEFAULT; + // timeout for one single network rpc private long apiCallAttemptTimeout = -1L; + // timeout for the whole write operation + private long timeout = Long.MAX_VALUE; private short bucketId; private boolean enableFastRetry; + // write context start private boolean retry; + private int retryCount; + private long requestTime = System.nanoTime(); + // write context end public WriteOptions throttleStrategy(ThrottleStrategy throttleStrategy) { this.throttleStrategy = throttleStrategy; @@ -150,6 +173,11 @@ public WriteOptions retry(boolean retry) { return this; } + public WriteOptions requestTime(long requestTime) { + this.requestTime = requestTime; + return this; + } + public ThrottleStrategy throttleStrategy() { return throttleStrategy; } @@ -162,7 +190,7 @@ public long apiCallAttemptTimeout() { return apiCallAttemptTimeout; } - // The value will be set by writer + // Writer will set the value WriteOptions bucketId(short bucketId) { this.bucketId = bucketId; return this; @@ -172,6 +200,11 @@ public short bucketId() { return bucketId; } + public WriteOptions timeout(long timeout) { + this.timeout = timeout; + return this; + } + public boolean enableFastRetry() { return enableFastRetry; } @@ -180,6 +213,24 @@ public boolean retry() { return retry; } + public int retryCountGetAndAdd() { + int oldRetryCount = this.retryCount; + this.retryCount = retryCount + 1; + return oldRetryCount; + } + + public int retryCount() { + return retryCount; + } + + public long requestTime() { + return requestTime; + } + + public long timeout() { + return timeout; + } + public WriteOptions copy() { WriteOptions copy = new WriteOptions(); copy.throttleStrategy = throttleStrategy; @@ -188,6 +239,9 @@ public WriteOptions copy() { copy.bucketId = bucketId; copy.enableFastRetry = enableFastRetry; copy.retry = retry; + copy.retryCount = retryCount; + copy.requestTime = requestTime; + copy.timeout = timeout; return copy; } } @@ -197,6 +251,7 @@ class ReadOptions { private ThrottleStrategy throttleStrategy = ThrottleStrategy.BYPASS; private short bucket = UNSET_BUCKET; + private int retryCount; public ReadOptions throttleStrategy(ThrottleStrategy throttleStrategy) { this.throttleStrategy = throttleStrategy; @@ -215,6 +270,16 @@ public ThrottleStrategy throttleStrategy() { public short bucket() { return bucket; } + + public int retryCountGetAndAdd() { + int oldRetryCount = this.retryCount; + this.retryCount = retryCount + 1; + return oldRetryCount; + } + + public int retryCount() { + return retryCount; + } } class WriteResult { diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageFactory.java b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageFactory.java index a36af23c27..09ba8c4829 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageFactory.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageFactory.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -16,19 +24,27 @@ import org.apache.commons.lang3.StringUtils; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; public class ObjectStorageFactory { + public static final String PROTOCOL_ROOT = "root"; + public static final String EXTENSION_TYPE_KEY = "type"; + public static final String EXTENSION_TYPE_MAIN = "main"; + public static final String EXTENSION_TYPE_BACKGROUND = "background"; private static volatile ObjectStorageFactory instance; private final Map> protocolHandlers = new HashMap<>(); static { + ObjectStorageFactory.instance().registerProtocolHandler(PROTOCOL_ROOT, builder -> { + throw new UnsupportedOperationException(); + }); ObjectStorageFactory.instance() .registerProtocolHandler("s3", builder -> AwsObjectStorage.builder() - .bucket(builder.bucketURI) + .bucket(builder.bucket) .tagging(builder.tagging) .inboundLimiter(builder.inboundLimiter) .outboundLimiter(builder.outboundLimiter) @@ -36,7 +52,11 @@ public class ObjectStorageFactory { .checkS3ApiModel(builder.checkS3ApiModel) .threadPrefix(builder.threadPrefix) .build()) - .registerProtocolHandler("mem", builder -> new MemoryObjectStorage(builder.bucketURI.bucketId())); + .registerProtocolHandler("mem", builder -> new MemoryObjectStorage(builder.bucket.bucketId())) + .registerProtocolHandler("file", builder -> + LocalFileObjectStorage.builder() + .bucket(builder.bucket) + .build()); } private ObjectStorageFactory() { @@ -52,6 +72,10 @@ public Builder builder(BucketURI bucket) { return new Builder().bucket(bucket); } + public Builder builder() { + return new Builder(); + } + public static ObjectStorageFactory instance() { if (instance == null) { synchronized (ObjectStorageFactory.class) { @@ -65,21 +89,35 @@ public static ObjectStorageFactory instance() { public class Builder { private final AtomicLong defaultThreadPrefixCounter = new AtomicLong(); - private BucketURI bucketURI; + private BucketURI bucket; + private List buckets; private Map tagging; private NetworkBandwidthLimiter inboundLimiter = NetworkBandwidthLimiter.NOOP; private NetworkBandwidthLimiter outboundLimiter = NetworkBandwidthLimiter.NOOP; private boolean readWriteIsolate; private boolean checkS3ApiModel = false; private String threadPrefix = ""; + private final Map extensions = new HashMap<>(); Builder bucket(BucketURI bucketURI) { - this.bucketURI = bucketURI; + this.bucket = bucketURI; return this; } public BucketURI bucket() { - return bucketURI; + return bucket; + } + + public Builder buckets(List buckets) { + this.buckets = buckets; + if (bucket == null && buckets.size() == 1) { + bucket = buckets.get(0); + } + return this; + } + + public List buckets() { + return buckets; } public Builder tagging(Map tagging) { @@ -128,6 +166,9 @@ public boolean checkS3ApiModel() { } public Builder threadPrefix(String prefix) { + if (prefix == null) { + return this; + } this.threadPrefix = prefix; return this; } @@ -136,11 +177,34 @@ public String threadPrefix() { return threadPrefix; } + public Builder extension(String key, Object value) { + this.extensions.put(key, value); + return this; + } + + @SuppressWarnings("unchecked") + public T extension(String key) { + return (T) this.extensions.get(key); + } + + public Map extensions() { + return extensions; + } + public ObjectStorage build() { if (StringUtils.isEmpty(this.threadPrefix)) { this.threadPrefix = Long.toString(defaultThreadPrefixCounter.getAndIncrement()); } - return protocolHandlers.get(bucketURI.protocol()).apply(this); + Function protocolHandle; + if (buckets != null && buckets.size() > 1) { + protocolHandle = protocolHandlers.get(PROTOCOL_ROOT); + } else { + protocolHandle = protocolHandlers.get(bucket.protocol()); + } + if (protocolHandle == null) { + throw new IllegalArgumentException("Cannot find protocol " + bucket.protocol()); + } + return protocolHandle.apply(this); } } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/ProxyWriter.java b/s3stream/src/main/java/com/automq/stream/s3/operator/ProxyWriter.java index 859b2405d2..e53a8f1709 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/ProxyWriter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/ProxyWriter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/RetryStrategy.java b/s3stream/src/main/java/com/automq/stream/s3/operator/RetryStrategy.java index 5bbbf33c24..2845e2c59f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/RetryStrategy.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/RetryStrategy.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/S3LatencyCalculator.java b/s3stream/src/main/java/com/automq/stream/s3/operator/S3LatencyCalculator.java index d99e4b3560..ae9399fcec 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/S3LatencyCalculator.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/S3LatencyCalculator.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficMonitor.java b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficMonitor.java new file mode 100644 index 0000000000..a17e6331b5 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficMonitor.java @@ -0,0 +1,59 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import com.automq.stream.s3.metrics.wrapper.Counter; + +/** + * A simple traffic monitor that counts the number of bytes transferred over time. + */ +class TrafficMonitor { + private final Counter trafficCounter = new Counter(); + private long lastTime = System.nanoTime(); + private long lastTraffic = 0; + + /** + * Record the number of bytes transferred. + * This method is thread-safe and can be called from multiple threads. + */ + public void record(long bytes) { + trafficCounter.inc(bytes); + } + + /** + * Calculate the rate of bytes transferred since the last call to this method, and reset the counter. + * Note: This method is not thread-safe. + */ + public double getRateAndReset() { + long currentTime = System.nanoTime(); + long deltaTime = currentTime - lastTime; + if (deltaTime <= 0) { + return 0; + } + + long currentTraffic = trafficCounter.get(); + long deltaTraffic = currentTraffic - lastTraffic; + + lastTime = currentTime; + lastTraffic = currentTraffic; + + return deltaTraffic * 1e9 / deltaTime; + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficRateLimiter.java b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficRateLimiter.java new file mode 100644 index 0000000000..eb9f127a00 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficRateLimiter.java @@ -0,0 +1,114 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import io.github.bucket4j.Bucket; +import io.github.bucket4j.BucketConfiguration; +import io.github.bucket4j.TokensInheritanceStrategy; + +/** + * A limiter that uses Bucker4j to limit the rate of network traffic. + */ +class TrafficRateLimiter { + + /** + * The maximum rate of refilling the Bucker4j bucket, which is 1 token per nanosecond. + */ + static final long MAX_BUCKET_TOKENS_PER_SECOND = TimeUnit.SECONDS.toNanos(1); + + /** + * The bucket used to limit the rate of network traffic in kilobytes per second. + * Maximum rate is 1 token per nanosecond, which is 1 TB/s and can be regarded as unlimited. + */ + private final Bucket bucket; + /** + * The scheduler used to schedule the rate limiting tasks. + * It should be shutdown outside of this class. + */ + private final ScheduledExecutorService scheduler; + + private long currentRate; + + /** + * Create a limiter without limiting (actually a limit of 1 TB/s). + */ + public TrafficRateLimiter(ScheduledExecutorService scheduler) { + this(scheduler, Long.MAX_VALUE); + } + + public TrafficRateLimiter(ScheduledExecutorService scheduler, long bytesPerSecond) { + this.currentRate = toKbps(bytesPerSecond); + this.bucket = Bucket.builder() + .addLimit(limit -> limit + .capacity(currentRate) + .refillGreedy(currentRate, Duration.ofSeconds(1)) + ).build(); + this.scheduler = scheduler; + } + + public long currentRate() { + return toBps(currentRate); + } + + /** + * Update the rate of the limiter. + * Note: The minimum rate is 1 KB/s and the maximum rate is 1 TB/s, any value outside this range will be + * clamped to this range. + * Note: An update will not take effect on the previous {@link this#consume} calls. For example, if the + * previous rate is 1 MB/s and the new rate is 10 MB/s, the previous {@link this#consume} calls will + * still be limited to 1 MB/s. + * Note：this method is not thread-safe. + */ + public void update(long bytesPerSecond) { + currentRate = toKbps(bytesPerSecond); + bucket.replaceConfiguration(BucketConfiguration.builder() + .addLimit(limit -> limit + .capacity(currentRate) + .refillGreedy(currentRate, Duration.ofSeconds(1)) + ).build(), + TokensInheritanceStrategy.PROPORTIONALLY + ); + } + + /** + * Consume the specified number of bytes and return a CompletableFuture that will be completed when the + * tokens are consumed. + * Note: DO NOT perform any heavy operations in the callback, otherwise it will block the scheduler. + */ + public CompletableFuture consume(long bytes) { + return bucket.asScheduler().consume(toKbps(bytes), scheduler); + } + + private static long toKbps(long bps) { + long kbps = bps >> 10; + kbps = Math.min(kbps, MAX_BUCKET_TOKENS_PER_SECOND); + kbps = Math.max(kbps, 1L); + return kbps; + } + + private static long toBps(long kbps) { + return kbps << 10; + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficRegulator.java b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficRegulator.java new file mode 100644 index 0000000000..0ebddfdf5d --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficRegulator.java @@ -0,0 +1,164 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import com.google.common.collect.Comparators; +import com.google.common.collect.EvictingQueue; +import com.google.common.math.Stats; + +import org.slf4j.Logger; + +import java.util.List; + +/** + * A traffic regulator that adjusts the rate of network traffic based on the success and failure rates. + */ +class TrafficRegulator { + private static final long MIN = 10L << 20; // 10 MB/s + private static final long MAX = TrafficRateLimiter.MAX_BUCKET_TOKENS_PER_SECOND * 1024; // 953.67 GB/s + + private static final int HISTORY_SIZE = 64; + private static final int TOP_SUCCESS_RATE_COUNT = 4; + private static final double FAST_INCREMENT_RATIO = 0.5; + private static final double SLOW_INCREMENT_RATIO = 0.05; + + private static final int WINDOW_SIZE = 5; // 5 seconds + + private final String operation; + private final Logger logger; + /** + * A queue to store the success rate (with a zero failure rate) history. + */ + private final EvictingQueue successRateQueue = EvictingQueue.create(HISTORY_SIZE); + private final TrafficMonitor success; + private final TrafficMonitor failure; + private final TrafficRateLimiter rateLimiter; + private final TrafficVolumeLimiter volumeLimiter; + + public TrafficRegulator(String operation, TrafficMonitor success, TrafficMonitor failure, + TrafficRateLimiter rateLimiter, TrafficVolumeLimiter volumeLimiter, Logger logger) { + this.operation = operation; + this.success = success; + this.failure = failure; + this.rateLimiter = rateLimiter; + this.volumeLimiter = volumeLimiter; + this.logger = logger; + } + + /** + * The maximum request size allowed by the regulator. + * Any request larger than this size should be downscaled to this size. + */ + public long maxRequestSize() { + return MIN * WINDOW_SIZE; + } + + public void regulate() { + double successRate = success.getRateAndReset(); + double failureRate = failure.getRateAndReset(); + double totalRate = successRate + failureRate; + + maybeRecord(successRate, failureRate); + long newRate = calculateNewRate(rateLimiter.currentRate(), successRate, failureRate, totalRate); + rateLimiter.update(newRate); + volumeLimiter.update(newRate * WINDOW_SIZE); + } + + private void maybeRecord(double successRate, double failureRate) { + if (successRate > 0 && failureRate <= 0) { + successRateQueue.add(successRate); + } + } + + private double meanOfTopSuccessRates() { + if (successRateQueue.isEmpty()) { + return MIN; + } + + // Reduce the sample count on warmup + int topCount = ceilDivide(successRateQueue.size() * TOP_SUCCESS_RATE_COUNT, HISTORY_SIZE); + List tops = successRateQueue.stream() + .collect(Comparators.greatest(topCount, Double::compareTo)); + assert !tops.isEmpty(); + return Stats.meanOf(tops); + } + + private static int ceilDivide(int dividend, int divisor) { + return (dividend + divisor - 1) / divisor; + } + + private long calculateNewRate(double currentLimit, double successRate, double failureRate, double totalRate) { + boolean isIncrease = totalRate <= 0 || failureRate <= 0; + long newRate = isIncrease ? increase(currentLimit) : decrease(successRate); + + if (MAX == newRate) { + // skip logging + return newRate; + } + String action = isIncrease ? "Increase" : "Decrease"; + logger.info("{} {} limit, current limit: {}, success rate: {}, failure rate: {}, new rate: {}", + action, operation, formatRate(currentLimit), formatRate(successRate), formatRate(failureRate), formatRate(newRate)); + return newRate; + } + + private long increase(double currentLimit) { + if (MAX == currentLimit) { + return MAX; + } + double historyRate = meanOfTopSuccessRates(); + if (currentLimit > historyRate * (1 + SLOW_INCREMENT_RATIO * 120)) { + // If the current limit is higher enough, which means there is and will be no throttling, + // so we can just increase the limit to the maximum. + logger.info("{} limit is high enough, current limit: {}, history rate: {}, new rate: {}", + operation, formatRate(currentLimit), formatRate(historyRate), formatRate(MAX)); + return MAX; + } + + List newRates = List.of( + currentLimit + historyRate * FAST_INCREMENT_RATIO, + currentLimit + historyRate * SLOW_INCREMENT_RATIO, + historyRate + ); + + // Find 2nd largest new rate + double newRate = newRates.stream() + .sorted() + .skip(1) + .findFirst() + .get(); + return (long) Math.min(newRate, MAX); + } + + private long decrease(double successRate) { + return (long) Math.max(successRate, MIN); + } + + private static String formatRate(double bytesPerSecond) { + String[] units = {" B/s", "KB/s", "MB/s", "GB/s", "TB/s"}; + int unitIndex = 0; + + while (bytesPerSecond >= 1024 && unitIndex < units.length - 1) { + bytesPerSecond /= 1024; + unitIndex++; + } + + return String.format("%6.2f %s", bytesPerSecond, units[unitIndex]); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficVolumeLimiter.java b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficVolumeLimiter.java new file mode 100644 index 0000000000..e242a8343f --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/TrafficVolumeLimiter.java @@ -0,0 +1,88 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import com.ibm.asyncutil.locks.AsyncSemaphore; +import com.ibm.asyncutil.locks.FairAsyncSemaphore; + +import java.util.concurrent.CompletableFuture; + +/** + * A limiter that uses an async semaphore to limit the volume of network traffic. + */ +public class TrafficVolumeLimiter { + + private static final long MAX_SEMAPHORE_PERMITS = FairAsyncSemaphore.MAX_PERMITS; + + /** + * The semaphore used to limit the volume of network traffic in bytes. + */ + private final AsyncSemaphore semaphore; + + /** + * The current volume of network traffic in bytes. + */ + private long currentVolume; + + /** + * Create a limiter without limiting. + */ + public TrafficVolumeLimiter() { + this(MAX_SEMAPHORE_PERMITS); + } + + public TrafficVolumeLimiter(long bytes) { + this.semaphore = new FairAsyncSemaphore(bytes); + this.currentVolume = bytes; + } + + public long currentVolume() { + return currentVolume; + } + + /** + * Update the current volume of network traffic. + * Note: this method is not thread-safe. + */ + public void update(long bytes) { + if (bytes > currentVolume) { + semaphore.release(bytes - currentVolume); + } else { + semaphore.acquire(currentVolume - bytes); + } + currentVolume = bytes; + } + + /** + * Consume the specified number of bytes and return a CompletableFuture that will be completed when the bytes are consumed. + * Note: DO NOT perform any heavy operations in the callback, otherwise it will block thread which calls {@link #release} + */ + public CompletableFuture acquire(long bytes) { + return semaphore.acquire(bytes).toCompletableFuture(); + } + + /** + * Release the specified number of bytes. + * It may complete a number of waiting futures returned by {@link #acquire} and execute their callbacks. + */ + public void release(long bytes) { + semaphore.release(bytes); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/Writer.java b/s3stream/src/main/java/com/automq/stream/s3/operator/Writer.java index a832e9addf..2884827811 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/Writer.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/Writer.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/WriterWrapper.java b/s3stream/src/main/java/com/automq/stream/s3/operator/WriterWrapper.java index 85dc4a6ca4..fa2a1c5180 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/WriterWrapper.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/WriterWrapper.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/main/java/com/automq/stream/s3/streams/StreamCloseHook.java b/s3stream/src/main/java/com/automq/stream/s3/streams/StreamCloseHook.java index 2f0b37aca6..6877fda766 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/streams/StreamCloseHook.java +++ b/s3stream/src/main/java/com/automq/stream/s3/streams/StreamCloseHook.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.streams; diff --git a/s3stream/src/main/java/com/automq/stream/s3/streams/StreamManager.java b/s3stream/src/main/java/com/automq/stream/s3/streams/StreamManager.java index 88088c328b..d2e3b594e6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/streams/StreamManager.java +++ b/s3stream/src/main/java/com/automq/stream/s3/streams/StreamManager.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.streams; @@ -36,6 +44,14 @@ public interface StreamManager { */ CompletableFuture> getStreams(List streamIds); + /** + * Add a stream metadata listener. + * @param streamId the stream id to listen + * @param listener {@link StreamMetadataListener} + * @return listener handle to remove listener + */ + StreamMetadataListener.Handle addMetadataListener(long streamId, StreamMetadataListener listener); + default CompletableFuture createStream() { return createStream(Collections.emptyMap()); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/streams/StreamMetadataListener.java b/s3stream/src/main/java/com/automq/stream/s3/streams/StreamMetadataListener.java new file mode 100644 index 0000000000..3d23754c52 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/streams/StreamMetadataListener.java @@ -0,0 +1,31 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.streams; + +import com.automq.stream.s3.metadata.StreamMetadata; + +public interface StreamMetadataListener { + + void onNewStreamMetadata(StreamMetadata metadata); + + interface Handle { + void close(); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java b/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java index b7ffb7e306..cd326112b6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java +++ b/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.trace; diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java b/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java index 500be0021e..c50be86692 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java +++ b/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.trace; diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java b/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java index d74be1b226..75218eb891 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java +++ b/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.trace; diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java b/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java index 1dead6deb6..c3827a55f6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java +++ b/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.trace; diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java b/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java index 910728c2b3..538846ef1d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java +++ b/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.trace.aop; diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java b/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java index f3982b9e68..c90155a622 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java +++ b/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.trace.context; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/AppendResult.java b/s3stream/src/main/java/com/automq/stream/s3/wal/AppendResult.java index 26c2d84478..bc9b50aa31 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/AppendResult.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/AppendResult.java @@ -1,27 +1,32 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal; -import java.util.concurrent.CompletableFuture; - public interface AppendResult { - // The pre-allocated starting offset of the record - long recordOffset(); - - CompletableFuture future(); - - interface CallbackResult { + /** + * The offset of the appended record. + */ + RecordOffset recordOffset(); - // The record before this offset (exclusive) has been flushed to disk - long flushedOffset(); - } + /** + * The offset that is valid for the next record. + */ + RecordOffset nextOffset(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/DefaultAppendResult.java b/s3stream/src/main/java/com/automq/stream/s3/wal/DefaultAppendResult.java new file mode 100644 index 0000000000..647f20259a --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/DefaultAppendResult.java @@ -0,0 +1,23 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +public record DefaultAppendResult(RecordOffset recordOffset, RecordOffset nextOffset) implements AppendResult { +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/DefaultWalHandle.java b/s3stream/src/main/java/com/automq/stream/s3/wal/DefaultWalHandle.java new file mode 100644 index 0000000000..10bd2e6a8a --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/DefaultWalHandle.java @@ -0,0 +1,72 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +import com.automq.stream.s3.operator.BucketURI; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.operator.ObjectStorageFactory; +import com.automq.stream.s3.wal.impl.object.ObjectReservationService; +import com.automq.stream.utils.IdURI; + +import java.util.Locale; +import java.util.concurrent.CompletableFuture; + +public class DefaultWalHandle implements WalHandle { + + private final String clusterId; + + public DefaultWalHandle(String clusterId) { + this.clusterId = clusterId; + } + + @Override + public CompletableFuture acquirePermission(int nodeId, long nodeEpoch, IdURI walConfig, + AcquirePermissionOptions options) { + //noinspection SwitchStatementWithTooFewBranches + switch (walConfig.protocol().toUpperCase(Locale.ENGLISH)) { + case "S3": { + return acquireObjectWALPermission(nodeId, nodeEpoch, walConfig, options); + } + default: { + throw new IllegalArgumentException(String.format("Unsupported WAL protocol %s in %s", walConfig.protocol(), walConfig)); + } + } + } + + @Override + public CompletableFuture releasePermission(IdURI walConfig, ReleasePermissionOptions options) { + //noinspection SwitchStatementWithTooFewBranches + switch (walConfig.protocol().toUpperCase(Locale.ENGLISH)) { + case "S3": { + return CompletableFuture.completedFuture(null); + } + default: { + throw new IllegalArgumentException(String.format("Unsupported WAL protocol %s in %s", walConfig.protocol(), walConfig)); + } + } + } + + private CompletableFuture acquireObjectWALPermission(int nodeId, long nodeEpoch, IdURI walConfig, + AcquirePermissionOptions options) { + ObjectStorage objectStorage = ObjectStorageFactory.instance().builder(BucketURI.parse(walConfig)).build(); + ObjectReservationService reservationService = new ObjectReservationService(clusterId, objectStorage, walConfig.id()); + return reservationService.acquire(nodeId, nodeEpoch, options.failoverMode()); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/OpenMode.java b/s3stream/src/main/java/com/automq/stream/s3/wal/OpenMode.java new file mode 100644 index 0000000000..18d4f849d3 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/OpenMode.java @@ -0,0 +1,26 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +public enum OpenMode { + READ_WRITE, + FAILOVER, + READ_ONLY, +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/RecordOffset.java b/s3stream/src/main/java/com/automq/stream/s3/wal/RecordOffset.java new file mode 100644 index 0000000000..fc36116c5a --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/RecordOffset.java @@ -0,0 +1,34 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +import io.netty.buffer.ByteBuf; + +public interface RecordOffset extends Comparable { + + ByteBuf buffer(); + + default byte[] bufferAsBytes() { + ByteBuf buffer = buffer().slice(); + byte[] bytes = new byte[buffer.readableBytes()]; + buffer.readBytes(bytes); + return bytes; + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/RecoverResult.java b/s3stream/src/main/java/com/automq/stream/s3/wal/RecoverResult.java index 18bb533798..8ae384bd2c 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/RecoverResult.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/RecoverResult.java @@ -1,23 +1,28 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal; -import io.netty.buffer.ByteBuf; +import com.automq.stream.s3.model.StreamRecordBatch; public interface RecoverResult { - ByteBuf record(); + StreamRecordBatch record(); - /** - * @see AppendResult#recordOffset() - */ - long recordOffset(); + RecordOffset recordOffset(); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/ReservationService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/ReservationService.java new file mode 100644 index 0000000000..289fc9bd23 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/ReservationService.java @@ -0,0 +1,58 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +import java.util.concurrent.CompletableFuture; + +public interface ReservationService { + ReservationService NOOP = new NoopReservationService(); + + /** + * Acquire the permission for the given node id and epoch. + * + * @param nodeId the current node id or the failed node id in failover mode + * @param epoch the current node epoch or the failed node epoch in failover mode + * @param failover whether this is in failover mode + * @return a future that will be completed when the permission is acquired + */ + CompletableFuture acquire(long nodeId, long epoch, boolean failover); + + /** + * Verify the permission for the given node id and epoch. + * + * @param nodeId the current node id or the failed node id in failover mode + * @param epoch the current node epoch or the failed node epoch in failover mode + * @param failover whether this is in failover mode + * @return a future that will be completed with the result of the verification + */ + CompletableFuture verify(long nodeId, long epoch, boolean failover); + + class NoopReservationService implements ReservationService { + @Override + public CompletableFuture acquire(long nodeId, long epoch, boolean failover) { + return CompletableFuture.completedFuture(null); + } + + @Override + public CompletableFuture verify(long nodeId, long epoch, boolean failover) { + return CompletableFuture.completedFuture(true); + } + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/WalFactory.java b/s3stream/src/main/java/com/automq/stream/s3/wal/WalFactory.java new file mode 100644 index 0000000000..d3e029856f --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/WalFactory.java @@ -0,0 +1,72 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +import com.automq.stream.utils.IdURI; + +public interface WalFactory { + + WriteAheadLog build(IdURI uri, BuildOptions options); + + + class BuildOptions { + private final long nodeEpoch; + private final OpenMode openMode; + + private BuildOptions(long nodeEpoch, OpenMode openMode) { + this.nodeEpoch = nodeEpoch; + this.openMode = openMode; + } + + public OpenMode openMode() { + return openMode; + } + + public long nodeEpoch() { + return nodeEpoch; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private long nodeEpoch; + private OpenMode openMode = OpenMode.READ_WRITE; + + public Builder openMode(OpenMode openMode) { + this.openMode = openMode; + return this; + } + + public Builder nodeEpoch(long nodeEpoch) { + this.nodeEpoch = nodeEpoch; + return this; + } + + public BuildOptions build() { + if (nodeEpoch <= 0) { + throw new IllegalArgumentException("The node epoch must be greater than 0"); + } + return new BuildOptions(nodeEpoch, openMode); + } + } + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/WalHandle.java b/s3stream/src/main/java/com/automq/stream/s3/wal/WalHandle.java new file mode 100644 index 0000000000..e0affa3daa --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/WalHandle.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal; + +import com.automq.stream.utils.IdURI; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +public interface WalHandle { + + /** + * Fence the specified WAL on the specified node and acquire permission to access it. + */ + CompletableFuture acquirePermission(int nodeId, long nodeEpoch, IdURI walConfig, AcquirePermissionOptions options); + + /** + * Release the permission to access the specified WAL. + */ + CompletableFuture releasePermission(IdURI walConfig, ReleasePermissionOptions options); + + class AcquirePermissionOptions { + private boolean failoverMode; + private long timeoutMs = TimeUnit.SECONDS.toMillis(20); + + public AcquirePermissionOptions failoverMode(boolean failoverMode) { + this.failoverMode = failoverMode; + return this; + } + + public boolean failoverMode() { + return failoverMode; + } + + public AcquirePermissionOptions timeoutMs(long timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + + public long timeoutMs() { + return timeoutMs; + } + } + + class ReleasePermissionOptions { + private long timeoutMs = TimeUnit.SECONDS.toMillis(20); + + public ReleasePermissionOptions timeoutMs(long timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + + public long timeoutMs() { + return timeoutMs; + } + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java b/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java index 0d8ada2c10..e4977a6ef1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java @@ -1,16 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal; +import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.trace.context.TraceContext; import com.automq.stream.s3.wal.common.WALMetadata; import com.automq.stream.s3.wal.exception.OverCapacityException; @@ -18,6 +27,7 @@ import java.io.IOException; import java.util.Iterator; +import java.util.List; import java.util.concurrent.CompletableFuture; import io.netty.buffer.ByteBuf; @@ -35,6 +45,11 @@ public interface WriteAheadLog { */ WALMetadata metadata(); + /** + * Get WAL config URI string. We could use the uri to reconstruct the WAL instance. + */ + String uri(); + /** * Append data to log, note append may be out of order. * ex. when sequence append R1 R2 , R2 maybe complete before R1. @@ -42,19 +57,14 @@ public interface WriteAheadLog { * * @return The data position will be written. */ - AppendResult append(TraceContext context, ByteBuf data, int crc) throws OverCapacityException; + // TODO: change the doc + CompletableFuture append(TraceContext context, StreamRecordBatch streamRecordBatch) throws OverCapacityException; - default AppendResult append(TraceContext context, ByteBuf data) throws OverCapacityException { - return append(context, data, 0); - } + CompletableFuture get(RecordOffset recordOffset); - default AppendResult append(ByteBuf data, int crc) throws OverCapacityException { - return append(TraceContext.DEFAULT, data, crc); - } + CompletableFuture> get(RecordOffset startOffset, RecordOffset endOffset); - default AppendResult append(ByteBuf data) throws OverCapacityException { - return append(TraceContext.DEFAULT, data, 0); - } + RecordOffset confirmOffset(); /** * Recover log from the beginning. The iterator will return the recovered result in order. @@ -78,5 +88,5 @@ default AppendResult append(ByteBuf data) throws OverCapacityException { * @param offset inclusive trim offset. * @return future complete when trim done. */ - CompletableFuture trim(long offset); + CompletableFuture trim(RecordOffset offset); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/BenchTool.java b/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/BenchTool.java deleted file mode 100644 index a2789fcc30..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/BenchTool.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.benchmark; - -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.wal.impl.block.BlockWALService; -import com.automq.stream.s3.wal.util.WALChannel; - -import net.sourceforge.argparse4j.inf.ArgumentParser; -import net.sourceforge.argparse4j.inf.ArgumentParserException; -import net.sourceforge.argparse4j.inf.Namespace; -import net.sourceforge.argparse4j.internal.HelpScreenException; - -import java.io.File; -import java.io.IOException; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.wal.util.WALUtil.isBlockDevice; - -public class BenchTool { - - public static Namespace parseArgs(ArgumentParser parser, String[] args) { - Namespace ns = null; - try { - ns = parser.parseArgs(args); - } catch (HelpScreenException e) { - System.exit(0); - } catch (ArgumentParserException e) { - parser.handleError(e); - System.exit(1); - } - return ns; - } - - public static void resetWALHeader(String path) throws IOException { - System.out.println("Resetting WAL header"); - if (isBlockDevice(path)) { - // block device - int capacity = BlockWALService.WAL_HEADER_TOTAL_CAPACITY; - WALChannel channel = WALChannel.builder(path).capacity(capacity).build(); - channel.open(); - ByteBuf buf = ByteBufAlloc.byteBuffer(capacity); - buf.writeZero(capacity); - channel.write(buf, 0); - buf.release(); - channel.close(); - } else { - // normal file - File file = new File(path); - if (file.isFile() && !file.delete()) { - throw new IOException("Failed to delete existing file " + file); - } - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/RecoverTool.java b/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/RecoverTool.java deleted file mode 100644 index 44470cb466..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/RecoverTool.java +++ /dev/null @@ -1,177 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.benchmark; - -import com.automq.stream.s3.StreamRecordBatchCodec; -import com.automq.stream.s3.model.StreamRecordBatch; -import com.automq.stream.s3.wal.RecoverResult; -import com.automq.stream.s3.wal.impl.block.BlockWALHeader; -import com.automq.stream.s3.wal.impl.block.BlockWALService; -import com.automq.stream.s3.wal.util.WALUtil; - -import net.sourceforge.argparse4j.ArgumentParsers; -import net.sourceforge.argparse4j.impl.Arguments; -import net.sourceforge.argparse4j.inf.ArgumentParser; -import net.sourceforge.argparse4j.inf.Namespace; - -import java.io.IOException; -import java.util.Iterator; -import java.util.function.Function; -import java.util.stream.StreamSupport; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.wal.benchmark.BenchTool.parseArgs; - -/** - * RecoverTool is a tool to recover records in a WAL manually. - * It extends {@link BlockWALService} to use tools provided by {@link BlockWALService} - */ -public class RecoverTool extends BlockWALService implements AutoCloseable { - - @SuppressWarnings("this-escape") - public RecoverTool(Config config) throws IOException { - super(BlockWALService.recoveryBuilder(config.path)); - super.start(); - } - - public static void main(String[] args) throws IOException { - Namespace ns = parseArgs(Config.parser(), args); - Config config = new Config(ns); - - try (RecoverTool tool = new RecoverTool(config)) { - tool.run(config); - } - } - - private void run(Config config) throws IOException { - BlockWALHeader header = super.tryReadWALHeader(); - System.out.println(header); - - Iterable recordsSupplier = () -> recover(header, config); - Function decoder = StreamRecordBatchCodec::decode; - Function stringer = decoder.andThen(StreamRecordBatch::toString); - Function offsetStringer = offset -> readableOffset(offset, header.getCapacity()); - StreamSupport.stream(recordsSupplier.spliterator(), false) - .map(it -> new RecoverResultWrapper(it, stringer, offsetStringer)) - .peek(System.out::println) - .forEach(RecoverResultWrapper::release); - } - - private Iterator recover(BlockWALHeader header, Config config) { - long recoverOffset = config.offset != null ? config.offset : header.getTrimOffset(); - long windowLength = config.windowLength != -1 ? config.windowLength : header.getSlidingWindowMaxLength(); - RecoverIterator iterator = new RecoverIterator(recoverOffset, windowLength, -1); - if (config.strict) { - iterator.strictMode(); - } - if (config.showInvalid) { - iterator.reportError(); - } - return iterator; - } - - private String readableOffset(long offset, long capacity) { - long physical = WALUtil.recordOffsetToPosition(offset, capacity, WAL_HEADER_TOTAL_CAPACITY); - long mod = physical % 4096; - return String.format("Offset{logical=%d, physical=%d, mod=%d}", offset, physical, mod); - } - - @Override - public void close() { - super.shutdownGracefully(); - } - - /** - * A wrapper for {@link RecoverResult} to provide a function to convert {@link RecoverResult#record} to string - */ - public static class RecoverResultWrapper { - private final RecoverResult inner; - /** - * A function to convert {@link RecoverResult#record} to string - */ - private final Function stringer; - private final Function offsetStringer; - - public RecoverResultWrapper(RecoverResult inner, Function stringer, Function offsetStringer) { - this.inner = inner; - this.stringer = stringer; - this.offsetStringer = offsetStringer; - } - - public void release() { - inner.record().release(); - } - - @Override - public String toString() { - String offset = offsetStringer.apply(inner.recordOffset()); - if (inner instanceof InvalidRecoverResult) { - InvalidRecoverResult invalid = (InvalidRecoverResult) inner; - return String.format("%s{", inner.getClass().getSimpleName()) - + "offset=" + offset - + ", error=" + invalid.detail() - + '}'; - } - return String.format("%s{", inner.getClass().getSimpleName()) - + "offset=" + offset - + String.format(", record=(%d)", inner.record().readableBytes()) + stringer.apply(inner.record()) - + '}'; - } - } - - public static class Config { - final String path; - final Long offset; - final Long windowLength; - final Boolean strict; - final Boolean showInvalid; - - Config(Namespace ns) { - this.path = ns.getString("path"); - this.offset = ns.getLong("offset"); - this.windowLength = ns.getLong("windowLength"); - this.strict = ns.getBoolean("strict"); - this.showInvalid = ns.getBoolean("showInvalid"); - } - - static ArgumentParser parser() { - ArgumentParser parser = ArgumentParsers - .newArgumentParser("RecoverTool") - .defaultHelp(true) - .description("Recover records in a WAL file"); - parser.addArgument("-p", "--path") - .required(true) - .help("Path of the WAL file"); - parser.addArgument("-o", "--offset") - .type(Long.class) - .help("Offset to start recovering, default to the trimmed offset in the WAL header"); - parser.addArgument("-w", "--window-length") - .dest("windowLength") - .type(Long.class) - .setDefault(-1L) - .help("Length of the sliding window, default to the value in the WAL header"); - parser.addArgument("-s", "--strict") - .type(Boolean.class) - .action(Arguments.storeTrue()) - .setDefault(false) - .help("Strict mode, which will stop when reaching the end of the window, default to false"); - parser.addArgument("-i", "--show-invalid") - .dest("showInvalid") - .type(Boolean.class) - .action(Arguments.storeTrue()) - .setDefault(false) - .help("Show invalid records, default to false"); - return parser; - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/RecoveryBench.java b/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/RecoveryBench.java deleted file mode 100644 index a903f7a0de..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/RecoveryBench.java +++ /dev/null @@ -1,149 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.benchmark; - -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.RecoverResult; -import com.automq.stream.s3.wal.WriteAheadLog; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.impl.block.BlockWALService; - -import net.sourceforge.argparse4j.ArgumentParsers; -import net.sourceforge.argparse4j.inf.ArgumentParser; -import net.sourceforge.argparse4j.inf.Namespace; - -import org.apache.commons.lang3.time.StopWatch; - -import java.io.IOException; -import java.util.Iterator; -import java.util.Random; -import java.util.concurrent.atomic.AtomicInteger; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.Unpooled; - -import static com.automq.stream.s3.wal.benchmark.BenchTool.parseArgs; -import static com.automq.stream.s3.wal.benchmark.BenchTool.resetWALHeader; - -/** - * RecoveryBench is a tool to benchmark the recovery performance of {@link BlockWALService} - */ -public class RecoveryBench implements AutoCloseable { - - private final WriteAheadLog log; - private Random random = new Random(); - - public RecoveryBench(Config config) throws IOException { - this.log = BlockWALService.builder(config.path, config.capacity).build().start(); - recoverAndReset(log); - } - - private static int recoverAndReset(WriteAheadLog wal) { - int recovered = 0; - for (Iterator it = wal.recover(); it.hasNext(); ) { - it.next().record().release(); - recovered++; - } - wal.reset().join(); - return recovered; - } - - public static void main(String[] args) throws Exception { - Namespace ns = parseArgs(Config.parser(), args); - Config config = new Config(ns); - - resetWALHeader(config.path); - try (RecoveryBench bench = new RecoveryBench(config)) { - bench.run(config); - } - } - - private void run(Config config) throws Exception { - writeRecords(config.numRecords, config.recordSizeBytes); - recoverRecords(config.path); - } - - private void writeRecords(int numRecords, int recordSizeBytes) throws OverCapacityException { - System.out.println("Writing " + numRecords + " records of size " + recordSizeBytes + " bytes"); - byte[] bytes = new byte[recordSizeBytes]; - random.nextBytes(bytes); - ByteBuf payload = Unpooled.wrappedBuffer(bytes).retain(); - - AtomicInteger appended = new AtomicInteger(); - for (int i = 0; i < numRecords; i++) { - AppendResult result = log.append(payload.retainedDuplicate()); - result.future().whenComplete((r, e) -> { - if (e != null) { - System.err.println("Failed to append record: " + e.getMessage()); - e.printStackTrace(); - } else { - appended.incrementAndGet(); - } - }); - } - System.out.println("Appended " + appended.get() + " records (may not be the final number)"); - } - - private void recoverRecords(String path) throws IOException { - System.out.println("Recovering records from " + path); - WriteAheadLog recoveryLog = BlockWALService.recoveryBuilder(path).build().start(); - StopWatch stopWatch = StopWatch.createStarted(); - int recovered = recoverAndReset(recoveryLog); - System.out.println("Recovered " + recovered + " records in " + stopWatch.getTime() + " ms"); - } - - @Override - public void close() { - log.shutdownGracefully(); - } - - static class Config { - // following fields are WAL configuration - final String path; - final Long capacity; - - // following fields are benchmark configuration - final Integer numRecords; - final Integer recordSizeBytes; - - Config(Namespace ns) { - this.path = ns.getString("path"); - this.capacity = ns.getLong("capacity"); - this.numRecords = ns.getInt("records"); - this.recordSizeBytes = ns.getInt("recordSize"); - } - - static ArgumentParser parser() { - ArgumentParser parser = ArgumentParsers - .newArgumentParser("RecoveryBench") - .defaultHelp(true) - .description("Benchmark the recovery performance of BlockWALService"); - parser.addArgument("-p", "--path") - .required(true) - .help("Path of the WAL file"); - parser.addArgument("-c", "--capacity") - .type(Long.class) - .setDefault((long) 3 << 30) - .help("Capacity of the WAL in bytes"); - parser.addArgument("--records") - .type(Integer.class) - .setDefault(1 << 20) - .help("number of records to write"); - parser.addArgument("--record-size") - .dest("recordSize") - .type(Integer.class) - .setDefault(1 << 10) - .help("size of each record in bytes"); - return parser; - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java b/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java deleted file mode 100644 index 931a074c4e..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java +++ /dev/null @@ -1,381 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.benchmark; - -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.WriteAheadLog; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.impl.block.BlockWALService; -import com.automq.stream.utils.ThreadUtils; -import com.automq.stream.utils.Threads; - -import net.sourceforge.argparse4j.ArgumentParsers; -import net.sourceforge.argparse4j.inf.ArgumentParser; -import net.sourceforge.argparse4j.inf.Namespace; - -import java.io.IOException; -import java.util.NavigableSet; -import java.util.Random; -import java.util.concurrent.ConcurrentSkipListSet; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.LockSupport; -import java.util.concurrent.locks.ReentrantLock; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.Unpooled; - -import static com.automq.stream.s3.wal.benchmark.BenchTool.parseArgs; -import static com.automq.stream.s3.wal.benchmark.BenchTool.resetWALHeader; - -/** - * WriteBench is a tool for benchmarking write performance of {@link BlockWALService} - */ -public class WriteBench implements AutoCloseable { - private static final int LOG_INTERVAL_SECONDS = 1; - private static final int TRIM_INTERVAL_MILLIS = 100; - - private final WriteAheadLog log; - private final TrimOffset trimOffset = new TrimOffset(); - - // Generate random payloads for this benchmark tool - private Random random = new Random(); - - public WriteBench(Config config) throws IOException { - BlockWALService.BlockWALServiceBuilder builder = BlockWALService.builder(config.path, config.capacity); - if (config.depth != null) { - builder.ioThreadNums(config.depth); - } - if (config.iops != null) { - builder.writeRateLimit(config.iops); - } - if (config.bandwidth != null) { - builder.writeBandwidthLimit(config.bandwidth); - } - this.log = builder.build(); - this.log.start(); - this.log.reset(); - } - - public static void main(String[] args) throws IOException { - Namespace ns = parseArgs(Config.parser(), args); - Config config = new Config(ns); - - resetWALHeader(config.path); - try (WriteBench bench = new WriteBench(config)) { - bench.run(config); - } - } - - private static Runnable logIt(Config config, Stat stat) { - ScheduledExecutorService statExecutor = Threads.newSingleThreadScheduledExecutor( - ThreadUtils.createThreadFactory("stat-thread-%d", true), null); - statExecutor.scheduleAtFixedRate(() -> { - Stat.Result result = stat.reset(); - if (0 != result.count()) { - System.out.printf("Append task | Append Rate %d msg/s %d KB/s | Avg Latency %.3f ms | Max Latency %.3f ms\n", - TimeUnit.SECONDS.toNanos(1) * result.count() / result.elapsedTimeNanos(), - TimeUnit.SECONDS.toNanos(1) * (result.count() * config.recordSizeBytes) / result.elapsedTimeNanos() / 1024, - (double) result.costNanos() / TimeUnit.MILLISECONDS.toNanos(1) / result.count(), - (double) result.maxCostNanos() / TimeUnit.MILLISECONDS.toNanos(1)); - } - }, LOG_INTERVAL_SECONDS, LOG_INTERVAL_SECONDS, TimeUnit.SECONDS); - return statExecutor::shutdownNow; - } - - private void run(Config config) { - System.out.println("Starting benchmark"); - - ExecutorService executor = Threads.newFixedThreadPool( - config.threads, ThreadUtils.createThreadFactory("append-thread-%d", false), null); - AppendTaskConfig appendTaskConfig = new AppendTaskConfig(config); - Stat stat = new Stat(); - Runnable stopTrim = runTrimTask(); - for (int i = 0; i < config.threads; i++) { - int index = i; - executor.submit(() -> { - try { - runAppendTask(index, appendTaskConfig, stat); - } catch (Exception e) { - System.err.printf("Append task %d failed, %s\n", index, e.getMessage()); - e.printStackTrace(); - } - }); - } - Runnable stopLog = logIt(config, stat); - - executor.shutdown(); - try { - if (!executor.awaitTermination(config.durationSeconds + 10, TimeUnit.SECONDS)) { - executor.shutdownNow(); - } - } catch (InterruptedException e) { - executor.shutdownNow(); - } - stopLog.run(); - stopTrim.run(); - - System.out.println("Benchmark finished"); - } - - private Runnable runTrimTask() { - ScheduledExecutorService trimExecutor = Threads.newSingleThreadScheduledExecutor( - ThreadUtils.createThreadFactory("trim-thread-%d", true), null); - trimExecutor.scheduleAtFixedRate(() -> { - try { - log.trim(trimOffset.get()); - } catch (Exception e) { - System.err.printf("Trim task failed, %s\n", e.getMessage()); - e.printStackTrace(); - } - }, TRIM_INTERVAL_MILLIS, TRIM_INTERVAL_MILLIS, TimeUnit.MILLISECONDS); - return trimExecutor::shutdownNow; - } - - private void runAppendTask(int index, AppendTaskConfig config, Stat stat) throws Exception { - System.out.printf("Append task %d started\n", index); - - byte[] bytes = new byte[config.recordSizeBytes]; - random.nextBytes(bytes); - ByteBuf payload = Unpooled.wrappedBuffer(bytes).retain(); - int intervalNanos = (int) TimeUnit.SECONDS.toNanos(1) / Math.max(1, config.throughputBytes / config.recordSizeBytes); - long lastAppendTimeNanos = System.nanoTime(); - long taskStartTimeMillis = System.currentTimeMillis(); - - while (true) { - while (true) { - long now = System.nanoTime(); - long elapsedNanos = now - lastAppendTimeNanos; - if (elapsedNanos >= intervalNanos) { - lastAppendTimeNanos += intervalNanos; - break; - } - LockSupport.parkNanos((intervalNanos - elapsedNanos) >> 2); - } - - long now = System.currentTimeMillis(); - if (now - taskStartTimeMillis > TimeUnit.SECONDS.toMillis(config.durationSeconds)) { - break; - } - - long appendStartTimeNanos = System.nanoTime(); - AppendResult result; - try { - result = log.append(payload.retainedDuplicate()); - } catch (OverCapacityException e) { - System.err.printf("Append task %d failed, retry it, %s\n", index, e.getMessage()); - continue; - } - trimOffset.appended(result.recordOffset()); - result.future().thenAccept(v -> { - long costNanosValue = System.nanoTime() - appendStartTimeNanos; - stat.update(costNanosValue); - trimOffset.flushed(v.flushedOffset()); - }).whenComplete((v, e) -> { - if (e != null) { - System.err.printf("Append task %d failed, %s\n", index, e.getMessage()); - e.printStackTrace(); - } - }); - } - - System.out.printf("Append task %d finished\n", index); - } - - @Override - public void close() { - log.shutdownGracefully(); - } - - static class Config { - // following fields are WAL configuration - final String path; - final Long capacity; - final Integer depth; - final Integer iops; - final Long bandwidth; - - // following fields are benchmark configuration - final Integer threads; - final Integer throughputBytes; - final Integer recordSizeBytes; - final Long durationSeconds; - - Config(Namespace ns) { - this.path = ns.getString("path"); - this.capacity = ns.getLong("capacity"); - this.depth = ns.getInt("depth"); - this.iops = ns.getInt("iops"); - this.bandwidth = ns.getLong("bandwidth"); - this.threads = ns.getInt("threads"); - this.throughputBytes = ns.getInt("throughput"); - this.recordSizeBytes = ns.getInt("recordSize"); - this.durationSeconds = ns.getLong("duration"); - } - - static ArgumentParser parser() { - ArgumentParser parser = ArgumentParsers - .newArgumentParser("WriteBench") - .defaultHelp(true) - .description("Benchmark write performance of BlockWALService"); - parser.addArgument("-p", "--path") - .required(true) - .help("Path of the WAL file"); - parser.addArgument("-c", "--capacity") - .type(Long.class) - .setDefault((long) 1 << 30) - .help("Capacity of the WAL in bytes"); - parser.addArgument("-d", "--depth") - .type(Integer.class) - .help("IO depth of the WAL"); - parser.addArgument("--iops") - .type(Integer.class) - .help("IOPS of the WAL"); - parser.addArgument("--bandwidth") - .type(Long.class) - .help("Bandwidth of the WAL in bytes per second"); - parser.addArgument("--threads") - .type(Integer.class) - .setDefault(1) - .help("Number of threads to use to write"); - parser.addArgument("-t", "--throughput") - .type(Integer.class) - .setDefault(1 << 20) - .help("Expected throughput in total in bytes per second"); - parser.addArgument("-s", "--record-size") - .dest("recordSize") - .type(Integer.class) - .setDefault(1 << 10) - .help("Size of each record in bytes"); - parser.addArgument("-D", "--duration") - .type(Long.class) - .setDefault(60L) - .help("Duration of the benchmark in seconds"); - return parser; - } - } - - static class AppendTaskConfig { - final int throughputBytes; - final int recordSizeBytes; - final long durationSeconds; - - AppendTaskConfig(Config config) { - this.throughputBytes = config.throughputBytes / config.threads; - this.recordSizeBytes = config.recordSizeBytes; - this.durationSeconds = config.durationSeconds; - } - } - - static class Stat { - final AtomicLong count = new AtomicLong(); - final AtomicLong costNanos = new AtomicLong(); - final AtomicLong maxCostNanos = new AtomicLong(); - long lastResetTimeNanos = System.nanoTime(); - - public void update(long costNanosValue) { - count.incrementAndGet(); - costNanos.addAndGet(costNanosValue); - maxCostNanos.accumulateAndGet(costNanosValue, Math::max); - } - - /** - * NOT thread-safe - */ - public Result reset() { - long countValue = count.getAndSet(0); - long costNanosValue = costNanos.getAndSet(0); - long maxCostNanosValue = maxCostNanos.getAndSet(0); - - long now = System.nanoTime(); - long elapsedTimeNanos = now - lastResetTimeNanos; - lastResetTimeNanos = now; - - return new Result(countValue, costNanosValue, maxCostNanosValue, elapsedTimeNanos); - } - - public static final class Result { - private final long count; - private final long costNanos; - private final long maxCostNanos; - private final long elapsedTimeNanos; - - public Result(long count, long costNanos, long maxCostNanos, long elapsedTimeNanos) { - this.count = count; - this.costNanos = costNanos; - this.maxCostNanos = maxCostNanos; - this.elapsedTimeNanos = elapsedTimeNanos; - } - - public long count() { - return count; - } - - public long costNanos() { - return costNanos; - } - - public long maxCostNanos() { - return maxCostNanos; - } - - public long elapsedTimeNanos() { - return elapsedTimeNanos; - } - } - } - - public static class TrimOffset { - private final Lock lock = new ReentrantLock(); - // Offsets of all data appended but not yet flushed to disk - private final NavigableSet appendedOffsets = new ConcurrentSkipListSet<>(); - // Offset before which all data has been flushed to disk - private long flushedOffset = -1; - // Offset at which all data has been flushed to disk - private long committedOffset = -1; - - public void appended(long offset) { - appendedOffsets.add(offset); - } - - public void flushed(long offset) { - lock.lock(); - try { - if (offset > flushedOffset) { - flushedOffset = offset; - Long lower = appendedOffsets.lower(flushedOffset); - if (lower != null) { - appendedOffsets.headSet(lower).clear(); - committedOffset = lower; - } - } - } finally { - lock.unlock(); - } - } - - /** - * @return the offset at which all data has been flushed to disk, or -1 if no data has been flushed to disk - */ - public long get() { - lock.lock(); - try { - return committedOffset; - } finally { - lock.unlock(); - } - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/AppendResultImpl.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/AppendResultImpl.java deleted file mode 100644 index 8904192653..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/AppendResultImpl.java +++ /dev/null @@ -1,61 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.common; - -import com.automq.stream.s3.wal.AppendResult; - -import java.util.Objects; -import java.util.concurrent.CompletableFuture; - -public final class AppendResultImpl implements AppendResult { - private final long recordOffset; - private final CompletableFuture future; - - public AppendResultImpl(long recordOffset, CompletableFuture future) { - this.recordOffset = recordOffset; - this.future = future; - } - - @Override - public String toString() { - return "AppendResultImpl{" + "recordOffset=" + recordOffset + '}'; - } - - @Override - public long recordOffset() { - return recordOffset; - } - - @Override - public CompletableFuture future() { - return future; - } - - @Override - public boolean equals(Object obj) { - if (obj == this) { - return true; - } - if (obj == null || obj.getClass() != this.getClass()) { - return false; - } - var that = (AppendResultImpl) obj; - return this.recordOffset == that.recordOffset && - Objects.equals(this.future, that.future); - } - - @Override - public int hashCode() { - return Objects.hash(recordOffset, future); - } - -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/Record.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/Record.java index b8034cce70..3773b241fa 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/Record.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/common/Record.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.common; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java index 7c5bfb6e05..fff8955ef1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecordHeader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.common; @@ -18,69 +26,65 @@ public class RecordHeader { public static final int RECORD_HEADER_SIZE = 4 + 4 + 8 + 4 + 4; public static final int RECORD_HEADER_WITHOUT_CRC_SIZE = RECORD_HEADER_SIZE - 4; - public static final int RECORD_HEADER_MAGIC_CODE = 0x87654321; - - private int magicCode0 = RECORD_HEADER_MAGIC_CODE; - private int recordBodyLength1; - private long recordBodyOffset2; - private int recordBodyCRC3; + public static final int RECORD_HEADER_DATA_MAGIC_CODE = 0x87654321; + /** + * Magic code for record header indicating that the record body is empty (used for padding). + */ + public static final int RECORD_HEADER_EMPTY_MAGIC_CODE = 0x76543210; + + private final int magicCode0; + private final int recordBodyLength1; + private final long recordBodyOffset2; + private final int recordBodyCRC3; private int recordHeaderCRC4; - public static RecordHeader unmarshal(ByteBuf byteBuf) { - RecordHeader recordHeader = new RecordHeader(); + public RecordHeader(long offset, int length, int crc) { + this.magicCode0 = RECORD_HEADER_DATA_MAGIC_CODE; + this.recordBodyLength1 = length; + this.recordBodyOffset2 = offset + RECORD_HEADER_SIZE; + this.recordBodyCRC3 = crc; + } + + public RecordHeader(long offset, int length) { + this.magicCode0 = RECORD_HEADER_EMPTY_MAGIC_CODE; + this.recordBodyLength1 = length; + this.recordBodyOffset2 = offset + RECORD_HEADER_SIZE; + this.recordBodyCRC3 = 0; + } + + public RecordHeader(ByteBuf byteBuf) { byteBuf.markReaderIndex(); - recordHeader.magicCode0 = byteBuf.readInt(); - recordHeader.recordBodyLength1 = byteBuf.readInt(); - recordHeader.recordBodyOffset2 = byteBuf.readLong(); - recordHeader.recordBodyCRC3 = byteBuf.readInt(); - recordHeader.recordHeaderCRC4 = byteBuf.readInt(); + this.magicCode0 = byteBuf.readInt(); + this.recordBodyLength1 = byteBuf.readInt(); + this.recordBodyOffset2 = byteBuf.readLong(); + this.recordBodyCRC3 = byteBuf.readInt(); + this.recordHeaderCRC4 = byteBuf.readInt(); byteBuf.resetReaderIndex(); - return recordHeader; } public int getMagicCode() { return magicCode0; } - public RecordHeader setMagicCode(int magicCode) { - this.magicCode0 = magicCode; - return this; - } - public int getRecordBodyLength() { return recordBodyLength1; } - public RecordHeader setRecordBodyLength(int recordBodyLength) { - this.recordBodyLength1 = recordBodyLength; - return this; - } - public long getRecordBodyOffset() { return recordBodyOffset2; } - public RecordHeader setRecordBodyOffset(long recordBodyOffset) { - this.recordBodyOffset2 = recordBodyOffset; - return this; - } - public int getRecordBodyCRC() { return recordBodyCRC3; } - public RecordHeader setRecordBodyCRC(int recordBodyCRC) { - this.recordBodyCRC3 = recordBodyCRC; - return this; - } - public int getRecordHeaderCRC() { return recordHeaderCRC4; } @Override public String toString() { - return "RecordHeaderCoreData{" + + return "RecordHeader{" + "magicCode=" + magicCode0 + ", recordBodyLength=" + recordBodyLength1 + ", recordBodyOffset=" + recordBodyOffset2 + @@ -97,16 +101,10 @@ private ByteBuf marshalHeaderExceptCRC(ByteBuf buf) { return buf; } - public ByteBuf marshal(ByteBuf emptyBuf, boolean calculateCRC) { + public ByteBuf marshal(ByteBuf emptyBuf) { assert emptyBuf.writableBytes() == RECORD_HEADER_SIZE; ByteBuf buf = marshalHeaderExceptCRC(emptyBuf); - - if (calculateCRC) { - buf.writeInt(WALUtil.crc32(buf, RECORD_HEADER_WITHOUT_CRC_SIZE)); - } else { - buf.writeInt(-1); - } - + buf.writeInt(WALUtil.crc32(buf, RECORD_HEADER_WITHOUT_CRC_SIZE)); return buf; } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecoverResultImpl.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecoverResultImpl.java index 44a45f1d7c..5266ac086f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecoverResultImpl.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/common/RecoverResultImpl.java @@ -1,27 +1,35 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.common; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.RecordOffset; import com.automq.stream.s3.wal.RecoverResult; import java.util.Objects; -import io.netty.buffer.ByteBuf; - public class RecoverResultImpl implements RecoverResult { - private final ByteBuf record; - private final long recordOffset; + private final StreamRecordBatch record; + private final RecordOffset recordOffset; - public RecoverResultImpl(ByteBuf record, long recordOffset) { + public RecoverResultImpl(StreamRecordBatch record, RecordOffset recordOffset) { this.record = record; this.recordOffset = recordOffset; } @@ -35,12 +43,12 @@ public String toString() { } @Override - public ByteBuf record() { + public StreamRecordBatch record() { return record; } @Override - public long recordOffset() { + public RecordOffset recordOffset() { return recordOffset; } @@ -53,8 +61,7 @@ public boolean equals(Object obj) { return false; } var that = (RecoverResultImpl) obj; - return Objects.equals(this.record, that.record) && - this.recordOffset == that.recordOffset; + return Objects.equals(this.record, that.record) && Objects.equals(this.recordOffset, that.recordOffset); } @Override diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/ShutdownType.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/ShutdownType.java index 92fb54bfa8..86ae9a03a1 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/ShutdownType.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/common/ShutdownType.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.common; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/common/WALMetadata.java b/s3stream/src/main/java/com/automq/stream/s3/wal/common/WALMetadata.java index f263c01284..640caf94d4 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/common/WALMetadata.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/common/WALMetadata.java @@ -1,16 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.common; + public class WALMetadata { private final int nodeId; private final long epoch; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/OverCapacityException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/OverCapacityException.java index 54931b6483..c95e0b27c3 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/OverCapacityException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/OverCapacityException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/RuntimeIOException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/RuntimeIOException.java index 9ea5800241..bd6254367c 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/RuntimeIOException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/RuntimeIOException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/UnmarshalException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/UnmarshalException.java index 5013ca5ae5..e38bd29520 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/UnmarshalException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/UnmarshalException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALCapacityMismatchException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALCapacityMismatchException.java index 85d4feb50a..84857f9b5e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALCapacityMismatchException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALCapacityMismatchException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALFencedException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALFencedException.java index 56b3d773c8..537db07105 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALFencedException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALFencedException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; @@ -14,6 +22,10 @@ import java.io.IOException; public class WALFencedException extends IOException { + + public WALFencedException() { + } + public WALFencedException(String message) { super(message); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALNotInitializedException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALNotInitializedException.java index 9710373f07..56b1421684 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALNotInitializedException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALNotInitializedException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALShutdownException.java b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALShutdownException.java index b5ac330ab9..da43eab5a0 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALShutdownException.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/exception/WALShutdownException.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.exception; diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/DefaultRecordOffset.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/DefaultRecordOffset.java new file mode 100644 index 0000000000..f51b91b4df --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/DefaultRecordOffset.java @@ -0,0 +1,110 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl; + +import com.automq.stream.s3.wal.RecordOffset; + +import java.util.Objects; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class DefaultRecordOffset implements RecordOffset { + private static final byte MAGIC = (byte) 0xA8; + private final long epoch; + private final long offset; + private final int size; + + private DefaultRecordOffset(long epoch, long offset, int size) { + this.epoch = epoch; + this.offset = offset; + this.size = size; + } + + public static DefaultRecordOffset of(long epoch, long recordOffset, int recordSize) { + return new DefaultRecordOffset(epoch, recordOffset, recordSize); + } + + public static DefaultRecordOffset of(ByteBuf buf) { + buf = buf.slice(); + byte magic = buf.readByte(); + if (magic != MAGIC) { + throw new IllegalArgumentException("Invalid magic: " + magic); + } + return new DefaultRecordOffset(buf.readLong(), buf.readLong(), buf.readInt()); + } + + public static DefaultRecordOffset of(RecordOffset recordOffset) { + if (recordOffset instanceof DefaultRecordOffset) { + return (DefaultRecordOffset) recordOffset; + } + return of(recordOffset.buffer()); + } + + public long epoch() { + return epoch; + } + + public long offset() { + return offset; + } + + public int size() { + return size; + } + + @Override + public ByteBuf buffer() { + ByteBuf buffer = Unpooled.buffer(1 + 8 + 8 + 4); + buffer.writeByte(MAGIC); + buffer.writeLong(epoch); + buffer.writeLong(this.offset); + buffer.writeInt(this.size); + return buffer; + } + + @Override + public String toString() { + return "DefaultRecordOffset{" + + "epoch=" + epoch + + ", offset=" + offset + + ", size=" + size + + '}'; + } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) + return false; + DefaultRecordOffset offset1 = (DefaultRecordOffset) o; + return epoch == offset1.epoch && offset == offset1.offset && size == offset1.size; + } + + @Override + public int hashCode() { + return Objects.hash(epoch, offset, size); + } + + @Override + public int compareTo(RecordOffset o) { + DefaultRecordOffset other = DefaultRecordOffset.of(o); + return Long.compare(this.offset, other.offset); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/MemoryWriteAheadLog.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/MemoryWriteAheadLog.java index d96be4474b..affbce2c3e 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/MemoryWriteAheadLog.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/MemoryWriteAheadLog.java @@ -1,18 +1,29 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.impl; +import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.trace.context.TraceContext; import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.DefaultAppendResult; +import com.automq.stream.s3.wal.RecordOffset; import com.automq.stream.s3.wal.RecoverResult; import com.automq.stream.s3.wal.WriteAheadLog; import com.automq.stream.s3.wal.common.RecordHeader; @@ -21,6 +32,7 @@ import java.io.IOException; import java.util.Iterator; +import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicLong; @@ -58,30 +70,46 @@ public WALMetadata metadata() { } @Override - public AppendResult append(TraceContext traceContext, ByteBuf data, int crc) throws OverCapacityException { + public String uri() { + return "0@mem://?"; + } + + @Override + public CompletableFuture append(TraceContext context, StreamRecordBatch streamRecordBatch) { if (full) { - data.release(); - throw new OverCapacityException("MemoryWriteAheadLog is full"); + streamRecordBatch.release(); + return CompletableFuture.failedFuture(new OverCapacityException("MemoryWriteAheadLog is full")); } - int dataLength = data.readableBytes(); + int dataLength = streamRecordBatch.encoded().readableBytes(); long offset = offsetAlloc.getAndAdd(RecordHeader.RECORD_HEADER_SIZE + dataLength); ByteBuf buffer = Unpooled.buffer(dataLength); - buffer.writeBytes(data); - data.release(); + buffer.writeBytes(streamRecordBatch.encoded()); + streamRecordBatch.release(); dataMap.put(offset, buffer); + return CompletableFuture.completedFuture(new DefaultAppendResult( + DefaultRecordOffset.of(0, offset, 0), + DefaultRecordOffset.of(0, offset + 1, 0) + )); + } + + @Override + public CompletableFuture get(RecordOffset recordOffset) { + return CompletableFuture.completedFuture(StreamRecordBatch.parse(dataMap.get(DefaultRecordOffset.of(recordOffset).offset()), false)); + } + + @Override + public CompletableFuture> get(RecordOffset startOffset, RecordOffset endOffset) { + List list = dataMap + .subMap(DefaultRecordOffset.of(startOffset).offset(), true, DefaultRecordOffset.of(endOffset).offset(), false) + .values().stream() + .map(buf -> StreamRecordBatch.parse(buf, false)).collect(Collectors.toList()); + return CompletableFuture.completedFuture(list); + } - return new AppendResult() { - @Override - public long recordOffset() { - return offset; - } - - @Override - public CompletableFuture future() { - return CompletableFuture.completedFuture(null); - } - }; + @Override + public RecordOffset confirmOffset() { + return DefaultRecordOffset.of(0, offsetAlloc.get(), 0); } @Override @@ -90,13 +118,13 @@ public Iterator recover() { .stream() .map(e -> (RecoverResult) new RecoverResult() { @Override - public ByteBuf record() { - return e.getValue(); + public StreamRecordBatch record() { + return StreamRecordBatch.parse(e.getValue(), false); } @Override - public long recordOffset() { - return e.getKey(); + public RecordOffset recordOffset() { + return DefaultRecordOffset.of(0, e.getKey(), 0); } }) .collect(Collectors.toList()) @@ -105,17 +133,18 @@ public long recordOffset() { @Override public CompletableFuture reset() { + dataMap.forEach((offset, buf) -> buf.release()); dataMap.clear(); return CompletableFuture.completedFuture(null); } @Override - public CompletableFuture trim(long offset) { - dataMap.headMap(offset) + public CompletableFuture trim(RecordOffset offset) { + dataMap.headMap(DefaultRecordOffset.of(offset).offset()) .forEach((key, value) -> { dataMap.remove(key); value.release(); }); return CompletableFuture.completedFuture(null); } -} +} \ No newline at end of file diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/Block.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/Block.java deleted file mode 100644 index de5cfe8391..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/Block.java +++ /dev/null @@ -1,92 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.common.Record; -import com.automq.stream.s3.wal.common.RecordHeader; -import com.automq.stream.s3.wal.util.WALUtil; - -import java.util.List; -import java.util.concurrent.CompletableFuture; - -import io.netty.buffer.ByteBuf; - -/** - * A Block contains multiple records, and will be written to the WAL in one batch. - */ -public interface Block { - /** - * The start offset of this block. - * Align to {@link WALUtil#BLOCK_SIZE} - */ - long startOffset(); - - /** - * Append a record to this block. - * Cannot be called after {@link #data()} is called. - * - * @param recordSize The size of this record. - * @param recordSupplier The supplier of this record. - * @param future The future of this record, which will be completed when the record is written to the WAL. - * @return The start offset of this record. If the size of this block exceeds the limit, return -1. - */ - long addRecord(long recordSize, RecordSupplier recordSupplier, - CompletableFuture future); - - /** - * Futures of all records in this block. - */ - List> futures(); - - default boolean isEmpty() { - return futures().isEmpty(); - } - - /** - * The content of this block, which contains multiple records. - * The first call of this method will marshal all records in this block to a ByteBuf. It will be cached for later calls. - */ - ByteBuf data(); - - /** - * The size of this block. - */ - long size(); - - /** - * The end offset of this block. - */ - default long endOffset() { - return startOffset() + size(); - } - - void release(); - - /** - * Called when this block is polled and sent to the writer. - * Used for metrics. - */ - void polled(); - - @FunctionalInterface - interface RecordSupplier { - /** - * Generate a record. - * - * @param recordStartOffset The start offset of this record. - * @param emptyHeader An empty {@link ByteBuf} with the size of {@link RecordHeader#RECORD_HEADER_SIZE}. It will be used to marshal the header. - * @return The record. - */ - Record get(long recordStartOffset, ByteBuf emptyHeader); - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockImpl.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockImpl.java deleted file mode 100644 index 61dc487396..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockImpl.java +++ /dev/null @@ -1,163 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.FixedSizeByteBufPool; -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.metrics.TimerUtil; -import com.automq.stream.s3.metrics.stats.StorageOperationStats; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.common.Record; -import com.automq.stream.s3.wal.util.WALUtil; -import com.automq.stream.utils.Systems; - -import java.util.LinkedList; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.TimeUnit; -import java.util.function.Supplier; -import java.util.stream.Collectors; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; - -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; - -public class BlockImpl implements Block { - - /** - * The pool for record headers. - */ - private static final FixedSizeByteBufPool HEADER_POOL = new FixedSizeByteBufPool(RECORD_HEADER_SIZE, 1024 * Systems.CPU_CORES); - - private final long startOffset; - /** - * The max size of this block. - * Any try to add a record to this block will fail if the size of this block exceeds this limit. - */ - private final long maxSize; - /** - * The soft limit of this block. - * Any try to add a record to this block will fail if the size of this block exceeds this limit, - * unless the block is empty. - */ - private final long softLimit; - private final List> futures = new LinkedList<>(); - private final List> recordSuppliers = new LinkedList<>(); - private final long startTime; - /** - * The next offset to write in this block. - * Align to {@link WALUtil#BLOCK_SIZE} - */ - private long nextOffset = 0; - /** - * Lazily generated records and data. - */ - private List records = null; - private CompositeByteBuf data = null; - - /** - * Create a block. - * {@link #release()} must be called when this block is no longer used. - */ - public BlockImpl(long startOffset, long maxSize, long softLimit) { - this.startOffset = startOffset; - this.maxSize = maxSize; - this.softLimit = softLimit; - this.startTime = System.nanoTime(); - } - - @Override - public long startOffset() { - return startOffset; - } - - /** - * Note: this method is NOT thread safe. - */ - @Override - public long addRecord(long recordSize, RecordSupplier recordSupplier, - CompletableFuture future) { - assert records == null; - long requiredCapacity = nextOffset + recordSize; - if (requiredCapacity > maxSize) { - return -1; - } - // if there is no record in this block, we can write a record larger than SOFT_BLOCK_SIZE_LIMIT - if (requiredCapacity > softLimit && !futures.isEmpty()) { - return -1; - } - - long recordOffset = startOffset + nextOffset; - recordSuppliers.add(() -> { - ByteBuf header = HEADER_POOL.get().retain(); - return recordSupplier.get(recordOffset, header); - }); - nextOffset += recordSize; - futures.add(future); - - return recordOffset; - } - - @Override - public List> futures() { - return futures; - } - - @Override - public ByteBuf data() { - maybeGenerateRecords(); - maybeGenerateData(); - return data; - } - - private void maybeGenerateRecords() { - if (null != records) { - return; - } - records = recordSuppliers.stream() - .map(Supplier::get) - .collect(Collectors.toUnmodifiableList()); - } - - private void maybeGenerateData() { - if (null != data) { - return; - } - data = ByteBufAlloc.compositeByteBuffer(); - for (Record record : records) { - data.addComponents(true, record.header(), record.body()); - } - } - - @Override - public long size() { - return nextOffset; - } - - @Override - public void release() { - if (null != data) { - data.release(); - } - if (null != records) { - records.stream() - .map(Record::header) - .forEach(HEADER_POOL::release); - } - } - - @Override - public void polled() { - StorageOperationStats.getInstance().appendWALBlockPolledStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockWALHeader.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockWALHeader.java deleted file mode 100644 index 94d87b3fd3..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockWALHeader.java +++ /dev/null @@ -1,207 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.wal.common.ShutdownType; -import com.automq.stream.s3.wal.exception.UnmarshalException; -import com.automq.stream.s3.wal.util.WALUtil; - -import java.util.concurrent.atomic.AtomicLong; - -import io.netty.buffer.ByteBuf; - -/** - *

    - * Layout: - *

    - * 0 - [4B] {@link BlockWALHeader#magicCode0} Magic code of the WAL header, used to verify the start of the WAL header - *

    - * 1 - [8B] {@link BlockWALHeader#capacity1} Capacity of the block device, which is configured by the application - * and should not be modified after the first start of the service - *

    - * 2 - [8B] {@link BlockWALHeader#trimOffset2} The logical start offset of the WAL, records before which are - * considered useless and have been deleted - *

    - * 3 - [8B] {@link BlockWALHeader#lastWriteTimestamp3} The timestamp of the last write to the WAL header, used to - * determine which WAL header is the latest when recovering - *

    - * 4 - [8B] {@link BlockWALHeader#slidingWindowMaxLength4} The maximum size of the sliding window, which can be - * scaled up when needed, and is used to determine when to stop recovering - *

    - * 5 - [4B] {@link BlockWALHeader#shutdownType5} The shutdown type of the service, {@link ShutdownType#GRACEFULLY} or - * {@link ShutdownType#UNGRACEFULLY} - *

    - * 6 - [4B] {@link BlockWALHeader#nodeId6} the node id of the WAL - *

    - * 7 - [4B] {@link BlockWALHeader#epoch7} the epoch id of the node - *

    - * 8 - [4B] {@link BlockWALHeader#crc8} CRC of the rest of the WAL header, used to verify the correctness of the - * WAL header - */ -public class BlockWALHeader { - public static final int WAL_HEADER_MAGIC_CODE = 0x12345678; - public static final int WAL_HEADER_SIZE = 4 // magic code - + 8 // capacity - + 8 // trim offset - + 8 // last write timestamp - + 8 // sliding window max length - + 4 // shutdown type - + 4 // node id - + 4 // node epoch - + 8; // crc - public static final int WAL_HEADER_WITHOUT_CRC_SIZE = WAL_HEADER_SIZE - 4; - private final AtomicLong trimOffset2 = new AtomicLong(-1); - private final AtomicLong flushedTrimOffset = new AtomicLong(0); - private final AtomicLong slidingWindowMaxLength4 = new AtomicLong(0); - private int magicCode0 = WAL_HEADER_MAGIC_CODE; - private long capacity1; - private long lastWriteTimestamp3 = System.nanoTime(); - private ShutdownType shutdownType5 = ShutdownType.UNGRACEFULLY; - private int nodeId6; - private long epoch7; - private int crc8; - - public BlockWALHeader(long capacity, long windowMaxLength) { - this.capacity1 = capacity; - this.slidingWindowMaxLength4.set(windowMaxLength); - } - - public static BlockWALHeader unmarshal(ByteBuf buf) throws UnmarshalException { - BlockWALHeader blockWalHeader = new BlockWALHeader(0, 0); - buf.markReaderIndex(); - blockWalHeader.magicCode0 = buf.readInt(); - blockWalHeader.capacity1 = buf.readLong(); - long trimOffset = buf.readLong(); - blockWalHeader.trimOffset2.set(trimOffset); - blockWalHeader.flushedTrimOffset.set(trimOffset); - blockWalHeader.lastWriteTimestamp3 = buf.readLong(); - blockWalHeader.slidingWindowMaxLength4.set(buf.readLong()); - blockWalHeader.shutdownType5 = ShutdownType.fromCode(buf.readInt()); - blockWalHeader.nodeId6 = buf.readInt(); - blockWalHeader.epoch7 = buf.readLong(); - blockWalHeader.crc8 = buf.readInt(); - buf.resetReaderIndex(); - - if (blockWalHeader.magicCode0 != WAL_HEADER_MAGIC_CODE) { - throw new UnmarshalException(String.format("WALHeader MagicCode not match, Recovered: [%d] expect: [%d]", blockWalHeader.magicCode0, WAL_HEADER_MAGIC_CODE)); - } - - int crc = WALUtil.crc32(buf, WAL_HEADER_WITHOUT_CRC_SIZE); - if (crc != blockWalHeader.crc8) { - throw new UnmarshalException(String.format("WALHeader CRC not match, Recovered: [%d] expect: [%d]", blockWalHeader.crc8, crc)); - } - - return blockWalHeader; - } - - public long getCapacity() { - return capacity1; - } - - public long getTrimOffset() { - return trimOffset2.get(); - } - - // Update the trim offset if the given trim offset is larger than the current one. - public BlockWALHeader updateTrimOffset(long trimOffset) { - trimOffset2.accumulateAndGet(trimOffset, Math::max); - return this; - } - - public long getFlushedTrimOffset() { - return flushedTrimOffset.get(); - } - - public void updateFlushedTrimOffset(long flushedTrimOffset) { - this.flushedTrimOffset.accumulateAndGet(flushedTrimOffset, Math::max); - } - - public long getLastWriteTimestamp() { - return lastWriteTimestamp3; - } - - public BlockWALHeader setLastWriteTimestamp(long lastWriteTimestamp) { - this.lastWriteTimestamp3 = lastWriteTimestamp; - return this; - } - - public long getSlidingWindowMaxLength() { - return slidingWindowMaxLength4.get(); - } - - public AtomicLong getAtomicSlidingWindowMaxLength() { - return slidingWindowMaxLength4; - } - - public ShutdownType getShutdownType() { - return shutdownType5; - } - - public BlockWALHeader setShutdownType(ShutdownType shutdownType) { - this.shutdownType5 = shutdownType; - return this; - } - - public int getNodeId() { - return nodeId6; - } - - public BlockWALHeader setNodeId(int nodeId) { - this.nodeId6 = nodeId; - return this; - } - - public long getEpoch() { - return epoch7; - } - - public BlockWALHeader setEpoch(long epoch) { - this.epoch7 = epoch; - return this; - } - - @Override - public String toString() { - return "WALHeader{" - + "magicCode=" + magicCode0 - + ", capacity=" + capacity1 - + ", trimOffset=" + trimOffset2 - + ", lastWriteTimestamp=" + lastWriteTimestamp3 - + ", slidingWindowMaxLength=" + slidingWindowMaxLength4 - + ", shutdownType=" + shutdownType5 - + ", nodeId=" + nodeId6 - + ", epoch=" + epoch7 - + ", crc=" + crc8 - + '}'; - } - - private ByteBuf marshalHeaderExceptCRC() { - ByteBuf buf = ByteBufAlloc.byteBuffer(WAL_HEADER_SIZE); - buf.writeInt(magicCode0); - buf.writeLong(capacity1); - buf.writeLong(trimOffset2.get()); - buf.writeLong(lastWriteTimestamp3); - buf.writeLong(slidingWindowMaxLength4.get()); - buf.writeInt(shutdownType5.getCode()); - buf.writeInt(nodeId6); - buf.writeLong(epoch7); - return buf; - } - - ByteBuf marshal() { - ByteBuf buf = marshalHeaderExceptCRC(); - this.crc8 = WALUtil.crc32(buf, WAL_HEADER_WITHOUT_CRC_SIZE); - buf.writeInt(crc8); - return buf; - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockWALService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockWALService.java deleted file mode 100644 index d9b721ab72..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/BlockWALService.java +++ /dev/null @@ -1,920 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.Config; -import com.automq.stream.s3.metrics.S3StreamMetricsManager; -import com.automq.stream.s3.metrics.TimerUtil; -import com.automq.stream.s3.metrics.stats.StorageOperationStats; -import com.automq.stream.s3.trace.TraceUtils; -import com.automq.stream.s3.trace.context.TraceContext; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.RecoverResult; -import com.automq.stream.s3.wal.WriteAheadLog; -import com.automq.stream.s3.wal.common.AppendResultImpl; -import com.automq.stream.s3.wal.common.RecordHeader; -import com.automq.stream.s3.wal.common.RecoverResultImpl; -import com.automq.stream.s3.wal.common.ShutdownType; -import com.automq.stream.s3.wal.common.WALMetadata; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.RuntimeIOException; -import com.automq.stream.s3.wal.exception.UnmarshalException; -import com.automq.stream.s3.wal.util.WALCachedChannel; -import com.automq.stream.s3.wal.util.WALChannel; -import com.automq.stream.s3.wal.util.WALUtil; -import com.automq.stream.utils.IdURI; -import com.automq.stream.utils.ThreadUtils; -import com.automq.stream.utils.Threads; - -import org.apache.commons.lang3.StringUtils; -import org.apache.commons.lang3.time.StopWatch; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.IOException; -import java.util.Collections; -import java.util.Iterator; -import java.util.NoSuchElementException; -import java.util.Objects; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.Lock; -import java.util.function.Function; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.Constants.CAPACITY_NOT_SET; -import static com.automq.stream.s3.Constants.NOOP_EPOCH; -import static com.automq.stream.s3.Constants.NOOP_NODE_ID; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_MAGIC_CODE; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_WITHOUT_CRC_SIZE; - -/** - * /** - * BlockWALService provides an infinite WAL, which is implemented based on block devices. - * The capacity of the block device is configured by the application and may be smaller than the system allocation. - *

    - * Usage: - *

    - * 1. Call {@link BlockWALService#start} to start the service. Any other methods will throw an - * {@link IllegalStateException} if called before {@link BlockWALService#start}. - *

    - * 2. Maybe call {@link BlockWALService#recover} to recover all untrimmed records if any. - *

    - * 3. Call {@link BlockWALService#reset} to reset the service. This will clear all records, so make sure - * all recovered records are processed before calling this method. - *

    - * 4. Call {@link BlockWALService#append} to append records. As records are written in a circular way similar to - * RingBuffer, if the caller does not call {@link BlockWALService#trim} in time, an {@link OverCapacityException} - * will be thrown when calling {@link BlockWALService#append}. - *

    - * 5. Call {@link BlockWALService#shutdownGracefully} to shut down the service gracefully, which will wait for - * all pending writes to complete. - *

    - * Implementation: - *

    - * WAL Header - *

    - * There are {@link BlockWALService#WAL_HEADER_COUNT} WAL headers, each of which is {@link WALUtil#BLOCK_SIZE} bytes. - * The WAL header is used to record the meta information of the WAL, and is used to recover the WAL when the service is restarted. - *

    - * Sliding Window - *

    - * The sliding window contains all records that have not been successfully written to the block device. - * So when recovering, we only need to try to recover the records in the sliding window. - *

    - * Record Header - *

    - * Layout: - *

    - * 0 - [4B] {@link RecordHeader#getMagicCode} Magic code of the record header, - * used to verify the start of the record header - *

    - * 1 - [4B] {@link RecordHeader#getRecordBodyLength} The length of the record body - *

    - * 2 - [8B] {@link RecordHeader#getRecordBodyOffset} The logical start offset of the record body - *

    - * 3 - [4B] {@link RecordHeader#getRecordBodyCRC} CRC of the record body, used to verify - * the correctness of the record body - *

    - * 4 - [4B] {@link RecordHeader#getRecordHeaderCRC} CRC of the rest of the record header, - * used to verify the correctness of the record header - */ -public class BlockWALService implements WriteAheadLog { - public static final int WAL_HEADER_COUNT = 2; - public static final int WAL_HEADER_CAPACITY = WALUtil.BLOCK_SIZE; - public static final int WAL_HEADER_TOTAL_CAPACITY = WAL_HEADER_CAPACITY * WAL_HEADER_COUNT; - private static final Logger LOGGER = LoggerFactory.getLogger(BlockWALService.class); - private final AtomicBoolean started = new AtomicBoolean(false); - private final AtomicBoolean resetFinished = new AtomicBoolean(false); - private final AtomicLong writeHeaderRoundTimes = new AtomicLong(0); - private final ExecutorService walHeaderFlusher = Threads.newFixedThreadPool(1, ThreadUtils.createThreadFactory("flush-wal-header-thread-%d", true), LOGGER); - private long initialWindowSize; - private WALCachedChannel walChannel; - private SlidingWindowService slidingWindowService; - private BlockWALHeader walHeader; - private boolean recoveryMode; - private boolean firstStart; - private int nodeId = NOOP_NODE_ID; - private long epoch = NOOP_EPOCH; - - private BlockWALService() { - } - - /** - * A protected constructor for testing purpose. - */ - protected BlockWALService(BlockWALServiceBuilder builder) { - BlockWALService that = builder.build(); - this.initialWindowSize = that.initialWindowSize; - this.walChannel = that.walChannel; - this.slidingWindowService = that.slidingWindowService; - this.walHeader = that.walHeader; - this.recoveryMode = that.recoveryMode; - this.nodeId = that.nodeId; - this.epoch = that.epoch; - } - - public static BlockWALServiceBuilder builder(String path, long capacity) { - return new BlockWALServiceBuilder(path, capacity); - } - - public static BlockWALServiceBuilder builder(IdURI uri) { - BlockWALService.BlockWALServiceBuilder builder = BlockWALService.builder(uri.path(), uri.extensionLong("capacity", 2147483648L)); - Optional.ofNullable(uri.extensionString("iops")).filter(StringUtils::isNumeric).ifPresent(v -> builder.writeRateLimit(Integer.parseInt(v))); - Optional.ofNullable(uri.extensionString("iodepth")).filter(StringUtils::isNumeric).ifPresent(v -> builder.ioThreadNums(Integer.parseInt(v))); - Optional.ofNullable(uri.extensionString("iobandwidth")).filter(StringUtils::isNumeric).ifPresent(v -> builder.writeBandwidthLimit(Long.parseLong(v))); - return builder; - } - - public static BlockWALServiceBuilder recoveryBuilder(String path) { - return new BlockWALServiceBuilder(path).recoveryMode(true); - } - - private void flushWALHeader(ShutdownType shutdownType) throws IOException { - walHeader.setShutdownType(shutdownType); - flushWALHeader(); - } - - private synchronized void flushWALHeader() throws IOException { - long position = writeHeaderRoundTimes.getAndIncrement() % WAL_HEADER_COUNT * WAL_HEADER_CAPACITY; - walHeader.setLastWriteTimestamp(System.nanoTime()); - long trimOffset = walHeader.getTrimOffset(); - ByteBuf buf = walHeader.marshal(); - this.walChannel.retryWrite(buf, position); - this.walChannel.retryFlush(); - buf.release(); - walHeader.updateFlushedTrimOffset(trimOffset); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("WAL header flushed, position: {}, header: {}", position, walHeader); - } - } - - /** - * Try to read a record at the given offset. - * The returned record should be released by the caller. - * - * @throws ReadRecordException if the record is not found or the record is corrupted - */ - private ByteBuf readRecord(long recoverStartOffset, - Function logicalToPhysical) throws IOException, ReadRecordException { - final ByteBuf recordHeader = ByteBufAlloc.byteBuffer(RECORD_HEADER_SIZE); - RecordHeader readRecordHeader; - try { - readRecordHeader = parseRecordHeader(recoverStartOffset, recordHeader, logicalToPhysical); - } finally { - recordHeader.release(); - } - - int recordBodyLength = readRecordHeader.getRecordBodyLength(); - ByteBuf recordBody = ByteBufAlloc.byteBuffer(recordBodyLength); - try { - parseRecordBody(recoverStartOffset, readRecordHeader, recordBody, logicalToPhysical); - } catch (Exception e) { - recordBody.release(); - throw e; - } - - return recordBody; - } - - private RecordHeader parseRecordHeader(long recoverStartOffset, ByteBuf recordHeader, - Function logicalToPhysical) throws IOException, ReadRecordException { - final long position = logicalToPhysical.apply(recoverStartOffset); - int read = walChannel.retryRead(recordHeader, position); - if (read != RECORD_HEADER_SIZE) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset), - String.format("failed to read record header: expected %d bytes, actual %d bytes, recoverStartOffset: %d", RECORD_HEADER_SIZE, read, recoverStartOffset) - ); - } - - RecordHeader readRecordHeader = RecordHeader.unmarshal(recordHeader); - if (readRecordHeader.getMagicCode() != RECORD_HEADER_MAGIC_CODE) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset), - String.format("magic code mismatch: expected %d, actual %d, recoverStartOffset: %d", RECORD_HEADER_MAGIC_CODE, readRecordHeader.getMagicCode(), recoverStartOffset) - ); - } - - int recordHeaderCRC = readRecordHeader.getRecordHeaderCRC(); - int calculatedRecordHeaderCRC = WALUtil.crc32(recordHeader, RECORD_HEADER_WITHOUT_CRC_SIZE); - if (recordHeaderCRC != calculatedRecordHeaderCRC) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset), - String.format("record header crc mismatch: expected %d, actual %d, recoverStartOffset: %d", calculatedRecordHeaderCRC, recordHeaderCRC, recoverStartOffset) - ); - } - - int recordBodyLength = readRecordHeader.getRecordBodyLength(); - if (recordBodyLength <= 0) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset), - String.format("invalid record body length: %d, recoverStartOffset: %d", recordBodyLength, recoverStartOffset) - ); - } - - long recordBodyOffset = readRecordHeader.getRecordBodyOffset(); - if (recordBodyOffset != recoverStartOffset + RECORD_HEADER_SIZE) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset), - String.format("invalid record body offset: expected %d, actual %d, recoverStartOffset: %d", recoverStartOffset + RECORD_HEADER_SIZE, recordBodyOffset, recoverStartOffset) - ); - } - return readRecordHeader; - } - - private void parseRecordBody(long recoverStartOffset, RecordHeader readRecordHeader, - ByteBuf recordBody, Function logicalToPhysical) throws IOException, ReadRecordException { - long recordBodyOffset = readRecordHeader.getRecordBodyOffset(); - int recordBodyLength = readRecordHeader.getRecordBodyLength(); - long position = logicalToPhysical.apply(recordBodyOffset); - int read = walChannel.retryRead(recordBody, position); - if (read != recordBodyLength) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset + RECORD_HEADER_SIZE + recordBodyLength), - String.format("failed to read record body: expected %d bytes, actual %d bytes, recoverStartOffset: %d", recordBodyLength, read, recoverStartOffset) - ); - } - - int recordBodyCRC = readRecordHeader.getRecordBodyCRC(); - int calculatedRecordBodyCRC = WALUtil.crc32(recordBody); - if (recordBodyCRC != calculatedRecordBodyCRC) { - throw new ReadRecordException( - WALUtil.alignNextBlock(recoverStartOffset + RECORD_HEADER_SIZE + recordBodyLength), - String.format("record body crc mismatch: expected %d, actual %d, recoverStartOffset: %d", calculatedRecordBodyCRC, recordBodyCRC, recoverStartOffset) - ); - } - } - - @Override - public WriteAheadLog start() throws IOException { - if (started.get()) { - LOGGER.warn("block WAL service already started"); - return this; - } - StopWatch stopWatch = StopWatch.createStarted(); - - walChannel.open(channel -> Optional.ofNullable(tryReadWALHeader(walChannel)) - .map(BlockWALHeader::getCapacity) - .orElse(null)); - - BlockWALHeader header = tryReadWALHeader(walChannel); - if (null == header) { - assert !recoveryMode; - header = newWALHeader(); - firstStart = true; - LOGGER.info("no available WALHeader, create a new one: {}", header); - } else { - LOGGER.info("read WALHeader from WAL: {}", header); - } - - header.setShutdownType(ShutdownType.UNGRACEFULLY); - walHeaderReady(header); - - started.set(true); - LOGGER.info("block WAL service started, cost: {} ms", stopWatch.getTime(TimeUnit.MILLISECONDS)); - return this; - } - - private void registerMetrics() { - S3StreamMetricsManager.registerDeltaWalOffsetSupplier(() -> { - try { - return this.getCurrentStartOffset(); - } catch (Exception e) { - LOGGER.error("failed to get current start offset", e); - return 0L; - } - }, () -> walHeader.getFlushedTrimOffset()); - } - - private long getCurrentStartOffset() { - Lock lock = slidingWindowService.getBlockLock(); - lock.lock(); - try { - Block block = slidingWindowService.getCurrentBlockLocked(); - return block.startOffset() + block.size(); - } finally { - lock.unlock(); - } - } - - /** - * Protected method for testing purpose. - */ - protected BlockWALHeader tryReadWALHeader() throws IOException { - return tryReadWALHeader(walChannel); - } - - /** - * Try to read the header from WAL, return the latest one. - */ - private BlockWALHeader tryReadWALHeader(WALChannel walChannel) throws IOException { - BlockWALHeader header = null; - for (int i = 0; i < WAL_HEADER_COUNT; i++) { - ByteBuf buf = ByteBufAlloc.byteBuffer(BlockWALHeader.WAL_HEADER_SIZE); - try { - int read = walChannel.retryRead(buf, i * WAL_HEADER_CAPACITY); - if (read != BlockWALHeader.WAL_HEADER_SIZE) { - continue; - } - BlockWALHeader tmpHeader = BlockWALHeader.unmarshal(buf); - if (header == null || header.getLastWriteTimestamp() < tmpHeader.getLastWriteTimestamp()) { - header = tmpHeader; - } - } catch (UnmarshalException ignored) { - // failed to parse WALHeader, ignore - } finally { - buf.release(); - } - } - return header; - } - - private BlockWALHeader newWALHeader() { - return new BlockWALHeader(walChannel.capacity(), initialWindowSize); - } - - private void walHeaderReady(BlockWALHeader header) throws IOException { - if (nodeId != NOOP_NODE_ID) { - header.setNodeId(nodeId); - header.setEpoch(epoch); - } - this.walHeader = header; - flushWALHeader(); - } - - @Override - public void shutdownGracefully() { - StopWatch stopWatch = StopWatch.createStarted(); - - if (!started.getAndSet(false)) { - LOGGER.warn("block WAL service already shutdown or not started yet"); - return; - } - walHeaderFlusher.shutdown(); - try { - if (!walHeaderFlusher.awaitTermination(5, TimeUnit.SECONDS)) { - walHeaderFlusher.shutdownNow(); - } - } catch (InterruptedException e) { - walHeaderFlusher.shutdownNow(); - } - - boolean gracefulShutdown = Optional.ofNullable(slidingWindowService) - .map(s -> s.shutdown(1, TimeUnit.DAYS)) - .orElse(true); - try { - flushWALHeader(gracefulShutdown ? ShutdownType.GRACEFULLY : ShutdownType.UNGRACEFULLY); - } catch (IOException ignored) { - // shutdown anyway - } - - walChannel.close(); - - LOGGER.info("block WAL service shutdown gracefully: {}, cost: {} ms", gracefulShutdown, stopWatch.getTime(TimeUnit.MILLISECONDS)); - } - - @Override - public WALMetadata metadata() { - checkStarted(); - return new WALMetadata(walHeader.getNodeId(), walHeader.getEpoch()); - } - - @Override - public AppendResult append(TraceContext context, ByteBuf buf, int crc) throws OverCapacityException { - // get current method name - TraceContext.Scope scope = TraceUtils.createAndStartSpan(context, "BlockWALService::append"); - final long startTime = System.nanoTime(); - try { - AppendResult result = append0(buf, crc); - result.future().whenComplete((nil, ex) -> TraceUtils.endSpan(scope, ex)); - return result; - } catch (Throwable t) { - if (t instanceof OverCapacityException) { - StorageOperationStats.getInstance().appendWALFullStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); - } - buf.release(); - TraceUtils.endSpan(scope, t); - throw t; - } - } - - private AppendResult append0(ByteBuf body, int crc) throws OverCapacityException { - final long startTime = System.nanoTime(); - checkStarted(); - checkWriteMode(); - checkResetFinished(); - - final long recordSize = RECORD_HEADER_SIZE + body.readableBytes(); - final CompletableFuture appendResultFuture = new CompletableFuture<>(); - long expectedWriteOffset; - - Lock lock = slidingWindowService.getBlockLock(); - lock.lock(); - try { - Block block = slidingWindowService.getCurrentBlockLocked(); - Block.RecordSupplier recordSupplier = (offset, header) -> WALUtil.generateRecord(body, header, crc, offset); - expectedWriteOffset = block.addRecord(recordSize, recordSupplier, appendResultFuture); - if (expectedWriteOffset < 0) { - // this block is full, create a new one - block = slidingWindowService.sealAndNewBlockLocked(block, recordSize, walHeader.getFlushedTrimOffset(), walHeader.getCapacity() - WAL_HEADER_TOTAL_CAPACITY); - expectedWriteOffset = block.addRecord(recordSize, recordSupplier, appendResultFuture); - } - } finally { - lock.unlock(); - } - slidingWindowService.tryWriteBlock(); - - final AppendResult appendResult = new AppendResultImpl(expectedWriteOffset, appendResultFuture); - appendResult.future().whenComplete((nil, ex) -> StorageOperationStats.getInstance().appendWALCompleteStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS))); - StorageOperationStats.getInstance().appendWALBeforeStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); - return appendResult; - } - - @Override - public Iterator recover() { - checkStarted(); - if (firstStart) { - return Collections.emptyIterator(); - } - - long trimmedOffset = walHeader.getTrimOffset(); - long recoverStartOffset = trimmedOffset; - if (recoverStartOffset < 0) { - recoverStartOffset = 0; - } - long windowLength = walHeader.getSlidingWindowMaxLength(); - return new RecoverIterator(recoverStartOffset, windowLength, trimmedOffset); - } - - @Override - public CompletableFuture reset() { - checkStarted(); - - long newStartOffset = WALUtil.alignLargeByBlockSize(walHeader.getTrimOffset() + walHeader.getCapacity()); - - if (!recoveryMode) { - // in recovery mode, no need to start sliding window service - slidingWindowService.start(walHeader.getAtomicSlidingWindowMaxLength(), newStartOffset); - } - LOGGER.info("reset sliding window to offset: {}", newStartOffset); - CompletableFuture cf = trim(newStartOffset - 1, true) - .thenRun(() -> resetFinished.set(true)); - - if (!recoveryMode) { - // Only register metrics when not in recovery mode - return cf.thenRun(this::registerMetrics); - } - return cf; - } - - @Override - public CompletableFuture trim(long offset) { - return trim(offset, false); - } - - private CompletableFuture trim(long offset, boolean internal) { - checkStarted(); - if (!internal) { - checkWriteMode(); - checkResetFinished(); - if (offset >= slidingWindowService.getWindowCoreData().getStartOffset()) { - throw new IllegalArgumentException("failed to trim: record at offset " + offset + " has not been flushed yet"); - } - } - - walHeader.updateTrimOffset(offset); - return CompletableFuture.runAsync(() -> { - try { - flushWALHeader(); - } catch (IOException e) { - throw new RuntimeIOException(e); - } - }, walHeaderFlusher); - } - - private void checkStarted() { - if (!started.get()) { - throw new IllegalStateException("WriteAheadLog has not been started yet"); - } - } - - private void checkWriteMode() { - if (recoveryMode) { - throw new IllegalStateException("WriteAheadLog is in recovery mode"); - } - } - - private void checkResetFinished() { - if (!resetFinished.get()) { - throw new IllegalStateException("WriteAheadLog has not been reset yet"); - } - } - - private SlidingWindowService.WALHeaderFlusher flusher() { - return () -> flushWALHeader(ShutdownType.UNGRACEFULLY); - } - - public static class BlockWALServiceBuilder { - private final String blockDevicePath; - private long blockDeviceCapacityWant = CAPACITY_NOT_SET; - private Boolean direct = null; - private int initBufferSize = 1 << 20; // 1MiB - private int maxBufferSize = 1 << 24; // 16MiB - private int ioThreadNums = 8; - private long slidingWindowInitialSize = 1 << 20; // 1MiB - private long slidingWindowUpperLimit = 1 << 29; // 512MiB - private long slidingWindowScaleUnit = 1 << 22; // 4MiB - private long blockSoftLimit = 1 << 18; // 256KiB - // wal io request limit - private int writeRateLimit = 3000; - // wal io bandwidth limit - private long writeBandwidthLimit = Long.MAX_VALUE; // no limitation - private int nodeId = NOOP_NODE_ID; - private long epoch = NOOP_EPOCH; - private boolean recoveryMode = false; - - public BlockWALServiceBuilder(String blockDevicePath, long capacity) { - this.blockDevicePath = blockDevicePath; - this.blockDeviceCapacityWant = capacity; - } - - public BlockWALServiceBuilder(String blockDevicePath) { - this.blockDevicePath = blockDevicePath; - } - - public BlockWALServiceBuilder recoveryMode(boolean recoveryMode) { - this.recoveryMode = recoveryMode; - return this; - } - - public BlockWALServiceBuilder capacity(long capacity) { - this.blockDeviceCapacityWant = capacity; - return this; - } - - public BlockWALServiceBuilder config(Config config) { - return this - .nodeId(config.nodeId()) - .epoch(config.nodeEpoch()); - } - - public BlockWALServiceBuilder direct(boolean direct) { - this.direct = direct; - return this; - } - - public BlockWALServiceBuilder initBufferSize(int initBufferSize) { - this.initBufferSize = initBufferSize; - return this; - } - - public BlockWALServiceBuilder maxBufferSize(int maxBufferSize) { - this.maxBufferSize = maxBufferSize; - return this; - } - - public BlockWALServiceBuilder ioThreadNums(int ioThreadNums) { - this.ioThreadNums = ioThreadNums; - return this; - } - - public BlockWALServiceBuilder slidingWindowInitialSize(long slidingWindowInitialSize) { - this.slidingWindowInitialSize = slidingWindowInitialSize; - return this; - } - - public BlockWALServiceBuilder slidingWindowUpperLimit(long slidingWindowUpperLimit) { - this.slidingWindowUpperLimit = slidingWindowUpperLimit; - return this; - } - - public BlockWALServiceBuilder slidingWindowScaleUnit(long slidingWindowScaleUnit) { - this.slidingWindowScaleUnit = slidingWindowScaleUnit; - return this; - } - - public BlockWALServiceBuilder blockSoftLimit(long blockSoftLimit) { - this.blockSoftLimit = blockSoftLimit; - return this; - } - - public BlockWALServiceBuilder writeRateLimit(int writeRateLimit) { - this.writeRateLimit = writeRateLimit; - return this; - } - - public BlockWALServiceBuilder writeBandwidthLimit(long writeBandwidthLimit) { - this.writeBandwidthLimit = writeBandwidthLimit; - return this; - } - - public BlockWALServiceBuilder nodeId(int nodeId) { - this.nodeId = nodeId; - return this; - } - - public BlockWALServiceBuilder epoch(long epoch) { - this.epoch = epoch; - return this; - } - - public BlockWALService build() { - if (recoveryMode) { - assert blockDeviceCapacityWant == CAPACITY_NOT_SET; - assert nodeId == NOOP_NODE_ID; - assert epoch == NOOP_EPOCH; - } else { - // make blockDeviceCapacityWant align to BLOCK_SIZE - blockDeviceCapacityWant = blockDeviceCapacityWant / WALUtil.BLOCK_SIZE * WALUtil.BLOCK_SIZE; - } - - BlockWALService blockWALService = new BlockWALService(); - - WALChannel.WALChannelBuilder walChannelBuilder = WALChannel.builder(blockDevicePath) - .capacity(blockDeviceCapacityWant) - .initBufferSize(initBufferSize) - .maxBufferSize(maxBufferSize) - .recoveryMode(recoveryMode); - if (direct != null) { - walChannelBuilder.direct(direct); - } - WALChannel channel = walChannelBuilder.build(); - blockWALService.walChannel = WALCachedChannel.of(channel); - if (!blockWALService.walChannel.useDirectIO()) { - LOGGER.warn("block wal not using direct IO"); - } - - if (!recoveryMode) { - // in recovery mode, no need to create sliding window service - // make sure window size is less than capacity - slidingWindowInitialSize = Math.min(slidingWindowInitialSize, blockDeviceCapacityWant - WAL_HEADER_TOTAL_CAPACITY); - slidingWindowUpperLimit = Math.min(slidingWindowUpperLimit, blockDeviceCapacityWant - WAL_HEADER_TOTAL_CAPACITY); - blockWALService.initialWindowSize = slidingWindowInitialSize; - blockWALService.slidingWindowService = new SlidingWindowService( - channel, - ioThreadNums, - slidingWindowUpperLimit, - slidingWindowScaleUnit, - blockSoftLimit, - // leave some buffer for other write operations, for example, flush WAL header caused by trim - Math.max(writeRateLimit - 20, writeRateLimit / 2), - writeBandwidthLimit, - blockWALService.flusher() - ); - } - - blockWALService.recoveryMode = recoveryMode; - - if (nodeId != NOOP_NODE_ID) { - blockWALService.nodeId = nodeId; - blockWALService.epoch = epoch; - } - - LOGGER.info("build BlockWALService: {}", this); - - return blockWALService; - } - - @Override - public String toString() { - return "BlockWALServiceBuilder{" - + "blockDevicePath='" + blockDevicePath - + ", blockDeviceCapacityWant=" + blockDeviceCapacityWant - + ", direct=" + direct - + ", initBufferSize=" + initBufferSize - + ", maxBufferSize=" + maxBufferSize - + ", ioThreadNums=" + ioThreadNums - + ", slidingWindowInitialSize=" + slidingWindowInitialSize - + ", slidingWindowUpperLimit=" + slidingWindowUpperLimit - + ", slidingWindowScaleUnit=" + slidingWindowScaleUnit - + ", blockSoftLimit=" + blockSoftLimit - + ", writeRateLimit=" + writeRateLimit - + ", writeBandwidthLimit=" + writeBandwidthLimit - + ", nodeId=" + nodeId - + ", epoch=" + epoch - + ", recoveryMode=" + recoveryMode - + '}'; - } - } - - /** - * Only used for testing purpose. - */ - protected static class InvalidRecoverResult extends RecoverResultImpl { - private final String detail; - - InvalidRecoverResult(long recordOffset, String detail) { - super(ByteBufAlloc.byteBuffer(0), recordOffset); - this.detail = detail; - } - - public String detail() { - return detail; - } - - @Override - public boolean equals(Object obj) { - if (obj == this) { - return true; - } - if (obj == null || obj.getClass() != this.getClass()) { - return false; - } - var that = (InvalidRecoverResult) obj; - return Objects.equals(this.detail, that.detail) && - super.equals(obj); - } - - @Override - public int hashCode() { - return Objects.hash(detail, super.hashCode()); - } - } - - static class ReadRecordException extends Exception { - long jumpNextRecoverOffset; - - public ReadRecordException(long offset, String message) { - super(message); - this.jumpNextRecoverOffset = offset; - } - - public long getJumpNextRecoverOffset() { - return jumpNextRecoverOffset; - } - } - - /** - * Protected for testing purpose. - */ - protected class RecoverIterator implements Iterator { - private final long windowLength; - private final long skipRecordAtOffset; - private long nextRecoverOffset; - private long maybeFirstInvalidCycle = -1; - private long maybeFirstInvalidOffset = -1; - private RecoverResult next; - private boolean strictMode = false; - private long lastValidOffset = -1; - private boolean reportError = false; - - public RecoverIterator(long nextRecoverOffset, long windowLength, long skipRecordAtOffset) { - this.nextRecoverOffset = nextRecoverOffset; - this.skipRecordAtOffset = skipRecordAtOffset; - this.windowLength = windowLength; - } - - /** - * Only used for testing purpose. - */ - public void strictMode() { - this.strictMode = true; - } - - /** - * Only used for testing purpose. - */ - public void reportError() { - this.reportError = true; - } - - @Override - public boolean hasNext() throws RuntimeIOException { - boolean hasNext = tryReadNextRecord(); - if (!hasNext) { - // recovery complete - walChannel.releaseCache(); - } - return hasNext; - } - - @Override - public RecoverResult next() throws RuntimeIOException { - if (!tryReadNextRecord()) { - throw new NoSuchElementException(); - } - - RecoverResult rst = next; - this.next = null; - return rst; - } - - /** - * Try to read next record. - * - * @return true if read success, false if no more record. {@link #next} will be null if and only if return false. - */ - private boolean tryReadNextRecord() throws RuntimeIOException { - if (next != null) { - return true; - } - while (shouldContinue()) { - long cycle = WALUtil.calculateCycle(nextRecoverOffset, walHeader.getCapacity(), WAL_HEADER_TOTAL_CAPACITY); - boolean skip = nextRecoverOffset == skipRecordAtOffset; - try { - ByteBuf nextRecordBody = readRecord(nextRecoverOffset, offset -> WALUtil.recordOffsetToPosition(offset, walHeader.getCapacity(), WAL_HEADER_TOTAL_CAPACITY)); - if (isOutOfWindow(nextRecoverOffset)) { - // should never happen, log it - LOGGER.error("[BUG] record offset out of window, offset: {}, firstInvalidOffset: {}, window: {}", - nextRecoverOffset, maybeFirstInvalidOffset, windowLength); - } - RecoverResultImpl recoverResult = new RecoverResultImpl(nextRecordBody, nextRecoverOffset); - lastValidOffset = nextRecoverOffset; - - nextRecoverOffset += RECORD_HEADER_SIZE + nextRecordBody.readableBytes(); - - if (maybeFirstInvalidCycle != -1 && maybeFirstInvalidCycle != cycle) { - // we meet a valid record in the next cycle, so the "invalid" record we met before is not really invalid - maybeFirstInvalidOffset = -1; - maybeFirstInvalidCycle = -1; - } - - if (skip) { - nextRecordBody.release(); - continue; - } - - next = recoverResult; - return true; - } catch (ReadRecordException e) { - if (maybeFirstInvalidOffset == -1 && WALUtil.isAligned(nextRecoverOffset) && !skip) { - maybeFirstInvalidCycle = cycle; - maybeFirstInvalidOffset = nextRecoverOffset; - // maybe the first invalid offset - LOGGER.info("maybe meet the first invalid offset during recovery. cycle: {}, offset: {}, window: {}, detail: '{}'", - maybeFirstInvalidCycle, maybeFirstInvalidOffset, windowLength, e.getMessage()); - } - - if (reportError) { - next = new InvalidRecoverResult(nextRecoverOffset, e.getMessage()); - } - nextRecoverOffset = e.getJumpNextRecoverOffset(); - if (reportError) { - return true; - } - } catch (IOException e) { - LOGGER.error("failed to read record at offset {}", nextRecoverOffset, e); - throw new RuntimeIOException(e); - } - } - return false; - } - - private boolean shouldContinue() { - if (!isOutOfWindow(nextRecoverOffset)) { - // within the window - return true; - } - if (strictMode) { - // not in the window, and in strict mode, so we should stop - return false; - } - // allow to try to recover a little more records (no more than 4MiB) - return nextRecoverOffset < lastValidOffset + Math.min(windowLength, 1 << 22); - } - - private boolean isOutOfWindow(long offset) { - if (maybeFirstInvalidOffset == -1) { - return false; - } - return offset >= maybeFirstInvalidOffset + windowLength; - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/SlidingWindowService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/SlidingWindowService.java deleted file mode 100644 index d9cbd4d6bf..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/block/SlidingWindowService.java +++ /dev/null @@ -1,553 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.metrics.TimerUtil; -import com.automq.stream.s3.metrics.stats.StorageOperationStats; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.WALShutdownException; -import com.automq.stream.s3.wal.util.WALChannel; -import com.automq.stream.s3.wal.util.WALUtil; -import com.automq.stream.utils.FutureUtil; -import com.automq.stream.utils.ThreadUtils; -import com.automq.stream.utils.Threads; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.IOException; -import java.time.Duration; -import java.util.Collection; -import java.util.LinkedList; -import java.util.List; -import java.util.PriorityQueue; -import java.util.Queue; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.ReentrantLock; - -import io.github.bucket4j.BlockingBucket; -import io.github.bucket4j.Bucket; -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.wal.impl.block.BlockWALService.WAL_HEADER_TOTAL_CAPACITY; - -/** - * The sliding window contains all records that have not been flushed to the disk yet. - * All records are written to the disk asynchronously by the AIO thread pool. - * When the sliding window is full, the current thread will be blocked until the sliding window is expanded. - * When the asynchronous write is completed, the start offset of the sliding window will be updated. - */ -public class SlidingWindowService { - private static final Logger LOGGER = LoggerFactory.getLogger(SlidingWindowService.class.getSimpleName()); - /** - * The minimum interval between two scheduled write operations. At most 1000 per second. - * - * @see this#pollBlockScheduler - */ - private static final long MIN_SCHEDULED_WRITE_INTERVAL_NANOS = TimeUnit.SECONDS.toNanos(1) / 1000; - /** - * The maximum rate of refilling the Bucker4j bucket, which is 1 token per nanosecond. - */ - private static final long MAX_BUCKET_TOKENS_PER_SECOND = TimeUnit.SECONDS.toNanos(1); - - /** - * Number of threads for writing blocks. - */ - private final int ioThreadNums; - - /** - * The upper limit of the sliding window. - */ - private final long upperLimit; - /** - * The unit to scale out the sliding window. - */ - private final long scaleUnit; - - /** - * The soft limit of a block. - * "Soft limit" means that the block can exceed this limit there is only one large record in this block. - */ - private final long blockSoftLimit; - - /** - * The rate limit of write operations. - */ - private final long writeRateLimit; - /** - * The bucket for rate limiting the write operations. - * - * @see #writeRateLimit - */ - private final Bucket writeRateBucket; - /** - * The bucket to limit the write bandwidth. - * Note: one token represents {@link WALUtil#BLOCK_SIZE} bytes. As the max rate in Bucket4j is 1 token per nanosecond, - * for a block size of 4KiB, the max rate is 3,814 GiB/s; for a block size of 512B, the max rate is 476 GiB/s. - */ - private final BlockingBucket writeBandwidthBucket; - - /** - * The channel to write data to the disk. - */ - private final WALChannel walChannel; - - /** - * The flusher used to flush the WAL header. - */ - private final WALHeaderFlusher walHeaderFlusher; - - /** - * The lock of {@link #pendingBlocks}, {@link #writingBlocks}, {@link #currentBlock}. - */ - private final Lock blockLock = new ReentrantLock(); - /** - * Blocks that are being written. - */ - private final Queue writingBlocks = new PriorityQueue<>(); - /** - * Whether the service is initialized. - * After the service is initialized, data in {@link #windowCoreData} is valid. - */ - private final AtomicBoolean initialized = new AtomicBoolean(false); - - /** - * The core data of the sliding window. Initialized when the service is started. - */ - private WindowCoreData windowCoreData; - /** - * Blocks that are waiting to be written. - * All blocks in this queue are ordered by the start offset. - */ - private volatile Queue pendingBlocks = new LinkedList<>(); - /** - * The current block, records are added to this block. - */ - private volatile Block currentBlock; - - /** - * The thread pool for write operations. - */ - private ExecutorService ioExecutor; - /** - * The scheduler for polling blocks and sending them to @{@link #ioExecutor}. - */ - private ScheduledExecutorService pollBlockScheduler; - - public SlidingWindowService(WALChannel walChannel, int ioThreadNums, long upperLimit, long scaleUnit, - long blockSoftLimit, int writeRateLimit, long writeBandwidthLimit, WALHeaderFlusher flusher) { - this.walChannel = walChannel; - this.ioThreadNums = ioThreadNums; - this.upperLimit = upperLimit; - this.scaleUnit = scaleUnit; - this.blockSoftLimit = blockSoftLimit; - this.writeRateLimit = writeRateLimit; - this.writeRateBucket = Bucket.builder() - .addLimit(limit -> limit - .capacity(Math.max(writeRateLimit / 10, 1)) - .refillGreedy(writeRateLimit, Duration.ofSeconds(1)) - ).build(); - long writeBandwidthLimitByBlock = Math.min(WALUtil.bytesToBlocks(writeBandwidthLimit), MAX_BUCKET_TOKENS_PER_SECOND); - this.writeBandwidthBucket = Bucket.builder() - .addLimit(limit -> limit - .capacity(Math.max(writeBandwidthLimitByBlock / 10, 1)) - .refillGreedy(writeBandwidthLimitByBlock, Duration.ofSeconds(1)) - ).build() - .asBlocking(); - this.walHeaderFlusher = flusher; - } - - public WindowCoreData getWindowCoreData() { - assert initialized(); - return windowCoreData; - } - - public void start(AtomicLong windowMaxLength, long windowStartOffset) { - this.windowCoreData = new WindowCoreData(windowMaxLength, windowStartOffset, windowStartOffset); - this.ioExecutor = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(ioThreadNums, - "block-wal-io-thread", false, LOGGER); - - long scheduledInterval = Math.max(MIN_SCHEDULED_WRITE_INTERVAL_NANOS, TimeUnit.SECONDS.toNanos(1) / writeRateLimit); - this.pollBlockScheduler = Threads.newSingleThreadScheduledExecutor( - ThreadUtils.createThreadFactory("wal-poll-block-thread-%d", false), LOGGER); - pollBlockScheduler.scheduleAtFixedRate(this::tryWriteBlock, 0, scheduledInterval, TimeUnit.NANOSECONDS); - - initialized.set(true); - } - - public boolean initialized() { - return initialized.get(); - } - - public boolean shutdown(long timeout, TimeUnit unit) { - if (this.ioExecutor == null) { - return true; - } - - boolean gracefulShutdown; - this.ioExecutor.shutdown(); - this.pollBlockScheduler.shutdownNow(); - List tasks = new LinkedList<>(); - try { - gracefulShutdown = this.ioExecutor.awaitTermination(timeout, unit); - } catch (InterruptedException e) { - tasks = this.ioExecutor.shutdownNow(); - gracefulShutdown = false; - } - - notifyWriteFuture(tasks); - - return gracefulShutdown; - } - - private void notifyWriteFuture(List tasks) { - Collection> futures = new LinkedList<>(); - for (Runnable task : tasks) { - if (task instanceof WriteBlockProcessor) { - WriteBlockProcessor processor = (WriteBlockProcessor) task; - futures.addAll(processor.block.futures()); - } - } - for (Block block : this.pendingBlocks) { - futures.addAll(block.futures()); - } - if (currentBlock != null && !currentBlock.isEmpty()) { - futures.addAll(currentBlock.futures()); - } - - doNotify(futures); - } - - private void doNotify(Collection> futures) { - for (CompletableFuture future : futures) { - future.completeExceptionally(new WALShutdownException("failed to write: ring buffer is shutdown")); - } - } - - /** - * Try to write a block. If it exceeds the rate limit, it will return immediately. - */ - public void tryWriteBlock() { - assert initialized(); - if (!tryAcquireWriteRateLimit()) { - return; - } - Block block = pollBlock(); - if (block != null) { - block.polled(); - ioExecutor.submit(new WriteBlockProcessor(block)); - } - } - - /** - * Try to acquire the write rate limit. - */ - private boolean tryAcquireWriteRateLimit() { - return writeRateBucket.tryConsume(1); - } - - public Lock getBlockLock() { - assert initialized(); - return blockLock; - } - - /** - * Seal and create a new block. It - * - puts the previous block to the write queue - * - creates a new block, sets it as the current block and returns it - * Note: this method is NOT thread safe, and it should be called with {@link #blockLock} locked. - */ - public Block sealAndNewBlockLocked(Block previousBlock, long minSize, long trimOffset, - long recordSectionCapacity) throws OverCapacityException { - assert initialized(); - long startOffset = nextBlockStartOffset(previousBlock); - - // If the end of the physical device is insufficient for this block, jump to the start of the physical device - if ((recordSectionCapacity - startOffset % recordSectionCapacity) < minSize) { - startOffset = startOffset + recordSectionCapacity - startOffset % recordSectionCapacity; - } - - // Not enough space for this block - if (startOffset + minSize - trimOffset > recordSectionCapacity) { - LOGGER.warn("failed to allocate write offset as the ring buffer is full: startOffset: {}, minSize: {}, trimOffset: {}, recordSectionCapacity: {}", - startOffset, minSize, trimOffset, recordSectionCapacity); - throw new OverCapacityException(String.format("failed to allocate write offset: ring buffer is full: startOffset: %d, minSize: %d, trimOffset: %d, recordSectionCapacity: %d", - startOffset, minSize, trimOffset, recordSectionCapacity)); - } - - long maxSize = upperLimit; - // The size of the block should not be larger than writable size of the ring buffer - // Let capacity=100, start=148, trim=49, then maxSize=100-148+49=1 - maxSize = Math.min(recordSectionCapacity - startOffset + trimOffset, maxSize); - // The size of the block should not be larger than the end of the physical device - // Let capacity=100, start=198, trim=198, then maxSize=100-198%100=2 - maxSize = Math.min(recordSectionCapacity - startOffset % recordSectionCapacity, maxSize); - - Block newBlock = new BlockImpl(startOffset, maxSize, blockSoftLimit); - if (!previousBlock.isEmpty()) { - // There are some records to be written in the previous block - pendingBlocks.add(previousBlock); - } else { - // The previous block is empty, so it can be released directly - previousBlock.release(); - } - setCurrentBlockLocked(newBlock); - return newBlock; - } - - /** - * Get the current block. - * Note: this method is NOT thread safe, and it should be called with {@link #blockLock} locked. - */ - public Block getCurrentBlockLocked() { - assert initialized(); - // The current block is null only when no record has been written - if (null == currentBlock) { - currentBlock = nextBlock(windowCoreData.getNextWriteOffset()); - } - return currentBlock; - } - - /** - * Set the current block. - * Note: this method is NOT thread safe, and it should be called with {@link #blockLock} locked. - */ - private void setCurrentBlockLocked(Block block) { - this.currentBlock = block; - } - - /** - * Get the start offset of the next block. - */ - private long nextBlockStartOffset(Block block) { - return block.startOffset() + WALUtil.alignLargeByBlockSize(block.size()); - } - - /** - * Create a new block with the given start offset. - * This method is only used when we don't know the maximum length of the new block. - */ - private Block nextBlock(long startOffset) { - // Trick: we cannot determine the maximum length of the block here, so we set it to 0 first. - // When we try to write a record, this block will be found full, and then a new block will be created. - return new BlockImpl(startOffset, 0, 0); - } - - /** - * Create a new block with the given previous block. - * This method is only used when we don't know the maximum length of the new block. - */ - private Block nextBlock(Block previousBlock) { - return nextBlock(nextBlockStartOffset(previousBlock)); - } - - /** - * Get a block to be written. If there is no non-empty block, return null. - */ - private Block pollBlock() { - blockLock.lock(); - try { - return pollBlockLocked(); - } finally { - blockLock.unlock(); - } - } - - /** - * Get a block to be written. If there is no non-empty block, return null. - * Note: this method is NOT thread safe, and it should be called with {@link #blockLock} locked. - */ - private Block pollBlockLocked() { - Block polled = null; - - Block currentBlock = getCurrentBlockLocked(); - if (!pendingBlocks.isEmpty()) { - polled = pendingBlocks.poll(); - } else if (currentBlock != null && !currentBlock.isEmpty()) { - polled = currentBlock; - setCurrentBlockLocked(nextBlock(currentBlock)); - } - - if (polled != null) { - writingBlocks.add(polled.startOffset()); - } - - return polled; - } - - /** - * Finish the given block, and return the start offset of the first block which has not been flushed yet. - */ - private long wroteBlock(Block wroteBlock) { - blockLock.lock(); - try { - return wroteBlockLocked(wroteBlock); - } finally { - blockLock.unlock(); - } - } - - /** - * Finish the given block, and return the start offset of the first block which has not been flushed yet. - * Note: this method is NOT thread safe, and it should be called with {@link #blockLock} locked. - */ - private long wroteBlockLocked(Block wroteBlock) { - boolean removed = writingBlocks.remove(wroteBlock.startOffset()); - assert removed; - if (writingBlocks.isEmpty()) { - return getCurrentBlockLocked().startOffset(); - } - return writingBlocks.peek(); - } - - private void writeBlockData(Block block) throws IOException { - final long start = System.nanoTime(); - long position = WALUtil.recordOffsetToPosition(block.startOffset(), walChannel.capacity(), WAL_HEADER_TOTAL_CAPACITY); - ByteBuf data = block.data(); - writeBandwidthBucket.consumeUninterruptibly(WALUtil.bytesToBlocks(data.readableBytes())); - walChannel.retryWrite(data, position); - walChannel.retryFlush(); - StorageOperationStats.getInstance().appendWALWriteStats.record(TimerUtil.timeElapsedSince(start, TimeUnit.NANOSECONDS)); - } - - private void makeWriteOffsetMatchWindow(long newWindowEndOffset) throws IOException { - // align to block size - newWindowEndOffset = WALUtil.alignLargeByBlockSize(newWindowEndOffset); - long windowStartOffset = windowCoreData.getStartOffset(); - long windowMaxLength = windowCoreData.getMaxLength(); - if (newWindowEndOffset > windowStartOffset + windowMaxLength) { - // endOffset - startOffset <= block.maxSize <= upperLimit in {@link #sealAndNewBlockLocked} - assert newWindowEndOffset - windowStartOffset <= upperLimit; - long newWindowMaxLength = Math.min(newWindowEndOffset - windowStartOffset + scaleUnit, upperLimit); - windowCoreData.scaleOutWindow(walHeaderFlusher, newWindowMaxLength); - } - } - - public interface WALHeaderFlusher { - void flush() throws IOException; - } - - public static class WindowCoreData { - private final Lock scaleOutLock = new ReentrantLock(); - private final AtomicLong maxLength; - /** - * Next write offset of sliding window, always aligned to the {@link WALUtil#BLOCK_SIZE}. - */ - private final AtomicLong nextWriteOffset; - /** - * Start offset of sliding window, always aligned to the {@link WALUtil#BLOCK_SIZE}. - * The data before this offset has already been written to the disk. - */ - private final AtomicLong startOffset; - - public WindowCoreData(AtomicLong maxLength, long nextWriteOffset, long startOffset) { - this.maxLength = maxLength; - this.nextWriteOffset = new AtomicLong(nextWriteOffset); - this.startOffset = new AtomicLong(startOffset); - } - - public long getMaxLength() { - return maxLength.get(); - } - - public void setMaxLength(long maxLength) { - this.maxLength.set(maxLength); - } - - public long getNextWriteOffset() { - return nextWriteOffset.get(); - } - - public long getStartOffset() { - return startOffset.get(); - } - - public void updateWindowStartOffset(long offset) { - this.startOffset.accumulateAndGet(offset, Math::max); - } - - public void scaleOutWindow(WALHeaderFlusher flusher, long newMaxLength) throws IOException { - boolean scaleWindowHappened = false; - scaleOutLock.lock(); - try { - if (newMaxLength < getMaxLength()) { - // Another thread has already scaled out the window. - return; - } - - setMaxLength(newMaxLength); - flusher.flush(); - scaleWindowHappened = true; - } finally { - scaleOutLock.unlock(); - if (scaleWindowHappened) { - LOGGER.info("window scale out to {}", newMaxLength); - } else { - LOGGER.debug("window already scale out, ignore"); - } - } - } - } - - class WriteBlockProcessor implements Runnable { - private final Block block; - private final long startTime; - - public WriteBlockProcessor(Block block) { - this.block = block; - this.startTime = System.nanoTime(); - } - - @Override - public void run() { - StorageOperationStats.getInstance().appendWALAwaitStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); - try { - writeBlock(this.block); - } catch (Exception e) { - // should not happen, but just in case - FutureUtil.completeExceptionally(block.futures().iterator(), e); - LOGGER.error(String.format("failed to write blocks, startOffset: %s", block.startOffset()), e); - } finally { - block.release(); - } - } - - private void writeBlock(Block block) throws IOException { - makeWriteOffsetMatchWindow(block.endOffset()); - writeBlockData(block); - - final long startTime = System.nanoTime(); - // Update the start offset of the sliding window after finishing writing the record. - windowCoreData.updateWindowStartOffset(wroteBlock(block)); - - FutureUtil.complete(block.futures().iterator(), new AppendResult.CallbackResult() { - @Override - public long flushedOffset() { - return windowCoreData.getStartOffset(); - } - - @Override - public String toString() { - return "CallbackResult{" + "flushedOffset=" + flushedOffset() + '}'; - } - }); - StorageOperationStats.getInstance().appendWALAfterStats.record(TimerUtil.timeElapsedSince(startTime, TimeUnit.NANOSECONDS)); - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java new file mode 100644 index 0000000000..15f224bec3 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java @@ -0,0 +1,309 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.exceptions.ObjectNotExistException; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.network.ThrottleStrategy; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; +import com.automq.stream.utils.Time; +import com.automq.stream.utils.threads.EventLoop; +import com.automq.stream.utils.threads.EventLoopSafe; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.NavigableMap; +import java.util.Queue; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentSkipListMap; + +import static com.automq.stream.s3.wal.impl.object.ObjectUtils.DATA_FILE_ALIGN_SIZE; +import static com.automq.stream.s3.wal.impl.object.ObjectUtils.floorAlignOffset; +import static com.automq.stream.s3.wal.impl.object.ObjectUtils.genObjectPathV1; + +@SuppressWarnings("checkstyle:cyclomaticComplexity") +@EventLoopSafe +public class DefaultReader { + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultReader.class); + private static final EventLoop[] EVENT_LOOPS = new EventLoop[4]; + + static { + for (int i = 0; i < EVENT_LOOPS.length; i++) { + EVENT_LOOPS[i] = new EventLoop("OBJECT_WAL_READER_" + i); + } + } + + private final ObjectStorage objectStorage; + private final String nodePrefix; + private final Time time; + + private final Queue singleReadTasks = new ConcurrentLinkedQueue<>(); + private final Queue batchReadTasks = new ConcurrentLinkedQueue<>(); + + // the rebuild task that is in process + private CompletableFuture rebuildIndexCf = CompletableFuture.completedFuture(null); + // the rebuild task that is waiting for sequential running. + private CompletableFuture awaitRebuildIndexCf; + private NavigableMap indexMap = new ConcurrentSkipListMap<>(); + private long indexLargestEpoch = -1L; + + private final EventLoop eventLoop; + + public DefaultReader(ObjectStorage objectStorage, String clusterId, int nodeId, String type, Time time) { + this.objectStorage = objectStorage; + this.nodePrefix = ObjectUtils.nodePrefix(clusterId, nodeId, type); + this.time = time; + this.eventLoop = EVENT_LOOPS[Math.abs(nodeId % EVENT_LOOPS.length)]; + } + + public CompletableFuture get(RecordOffset recordOffset) { + DefaultRecordOffset offset = recordOffset instanceof DefaultRecordOffset ? (DefaultRecordOffset) recordOffset : DefaultRecordOffset.of(recordOffset.buffer()); + SingleReadTask readTask = new SingleReadTask(offset.epoch(), offset.offset(), offset.size()); + singleReadTasks.add(readTask); + eventLoop.execute(this::doRunSingleGet); + return readTask.cf; + } + + public CompletableFuture> get(RecordOffset startOffset, RecordOffset endOffset) { + BatchReadTask readTask = new BatchReadTask(startOffset, endOffset); + if (readTask.startOffset.offset() == readTask.endOffset.offset()) { + return CompletableFuture.completedFuture(Collections.emptyList()); + } + batchReadTasks.add(readTask); + eventLoop.execute(this::doRunBatchGet); + return readTask.cf; + } + + private void doRunSingleGet() { + for (; ; ) { + SingleReadTask readTask = singleReadTasks.poll(); + if (readTask == null) { + break; + } + long objectStartObject = floorAlignOffset(readTask.offset); + String objectPath = genObjectPathV1(nodePrefix, readTask.epoch, objectStartObject); + long relativeStartOffset = readTask.offset - objectStartObject + WALObjectHeader.WAL_HEADER_SIZE_V1; + objectStorage.rangeRead( + new ObjectStorage.ReadOptions().bucket(objectStorage.bucketId()).throttleStrategy(ThrottleStrategy.BYPASS), + objectPath, + relativeStartOffset, + relativeStartOffset + readTask.size + ).whenCompleteAsync((buf, ex) -> { + try { + if (ex != null) { + readTask.cf.completeExceptionally(ex); + } else { + readTask.cf.complete(ObjectUtils.decodeRecordBuf(buf.slice())); + buf.release(); + } + } catch (Throwable e) { + readTask.cf.completeExceptionally(e); + } + }, eventLoop); + } + } + + private void doRunBatchGet() { + for (; ; ) { + BatchReadTask readTask = batchReadTasks.poll(); + if (readTask == null) { + break; + } + try { + doRunBatchGet0(readTask); + } catch (Throwable e) { + LOGGER.error("[UNEXPECTED] Failed to run {}", readTask, e); + readTask.cf.completeExceptionally(e); + } + } + } + + @SuppressWarnings("NPathComplexity") + private void doRunBatchGet0(BatchReadTask readTask) { + CompletableFuture indexCf; + if (indexLargestEpoch < readTask.endOffset.epoch()) { + indexCf = rebuildIndexMap(); + } else { + indexCf = CompletableFuture.completedFuture(null); + } + indexCf + .thenComposeAsync(nil -> { + NavigableMap indexMap; + Long floorKey = this.indexMap.floorKey(readTask.startOffset.offset()); + if (floorKey == null) { + indexMap = new TreeMap<>(); + } else { + indexMap = this.indexMap.subMap(floorKey, true, readTask.endOffset.offset(), false); + } + if (indexMap.isEmpty()) { + readTask.cf.completeExceptionally(new ObjectNotExistException( + String.format("Cannot find epoch for [%s, %s)", readTask.startOffset, readTask.endOffset) + )); + } + + List>> getCfList = new ArrayList<>(); + long nextGetOffset = readTask.startOffset.offset(); + List> entries = new ArrayList<>(indexMap.entrySet()); + for (int i = 0; i < entries.size(); i++) { + long epoch = entries.get(i).getValue(); + long epochEndOffset = (i == entries.size() - 1) ? Long.MAX_VALUE : entries.get(i + 1).getKey(); + while (nextGetOffset < epochEndOffset && nextGetOffset < readTask.endOffset.offset()) { + long objectStartOffset = floorAlignOffset(nextGetOffset); + String objectPath = genObjectPathV1(nodePrefix, epoch, objectStartOffset); + long relativeStartOffset = nextGetOffset - objectStartOffset + WALObjectHeader.WAL_HEADER_SIZE_V1; + // read to end + long finalNextGetOffset = nextGetOffset; + getCfList.add(objectStorage.rangeRead( + new ObjectStorage.ReadOptions().bucket(objectStorage.bucketId()).throttleStrategy(ThrottleStrategy.BYPASS), + objectPath, + relativeStartOffset, + -1 + ).thenApply(buf -> { + try { + List batches = new ArrayList<>(); + buf = buf.slice(); + long nextRecordOffset = finalNextGetOffset; + int lastReadableBytes = buf.readableBytes(); + while (buf.readableBytes() > 0 && nextRecordOffset < readTask.endOffset.offset()) { + StreamRecordBatch batch = ObjectUtils.decodeRecordBuf(buf); + boolean isTriggerTrimRecord = batch.getCount() == 0 && batch.getStreamId() == -1L && batch.getEpoch() == -1L; + if (!isTriggerTrimRecord) { + batches.add(batch); + } else { + batch.release(); + } + nextRecordOffset += lastReadableBytes - buf.readableBytes(); + lastReadableBytes = buf.readableBytes(); + } + return batches; + } finally { + buf.release(); + } + })); + nextGetOffset = objectStartOffset + DATA_FILE_ALIGN_SIZE; + } + } + return CompletableFuture.allOf(getCfList.toArray(new CompletableFuture[0])).whenCompleteAsync((nil2, ex) -> { + if (ex != null) { + getCfList.forEach(cf -> cf.thenAccept(l -> l.forEach(StreamRecordBatch::release))); + readTask.cf.completeExceptionally(ex); + return; + } + List batches = new ArrayList<>(); + for (CompletableFuture> cf : getCfList) { + batches.addAll(cf.join()); + } + readTask.cf.complete(batches); + }, eventLoop); + }, eventLoop) + .whenComplete((nil, ex) -> { + if (ex != null && !readTask.cf.isDone()) { + LOGGER.error("[UNEXPECTED] Failed to run {}", readTask, ex); + readTask.cf.completeExceptionally(ex); + } + }); + } + + private CompletableFuture rebuildIndexMap() { + if (rebuildIndexCf.isDone()) { + rebuildIndexCf = rebuildIndexMap0(); + return rebuildIndexCf; + } else { + if (awaitRebuildIndexCf != null) { + return awaitRebuildIndexCf; + } + awaitRebuildIndexCf = new CompletableFuture<>(); + CompletableFuture retCf = awaitRebuildIndexCf; + rebuildIndexCf.whenCompleteAsync((nil, ex) -> { + awaitRebuildIndexCf = null; + rebuildIndexCf = rebuildIndexMap0(); + rebuildIndexCf.whenComplete((nil2, ex2) -> { + if (ex2 != null) { + retCf.completeExceptionally(ex2); + } else { + retCf.complete(null); + } + }); + }, eventLoop); + return retCf; + } + } + + private CompletableFuture rebuildIndexMap0() { + return objectStorage.list(nodePrefix).thenAcceptAsync(list -> { + List objects = ObjectUtils.parse(list); + TreeMap newIndexMap = new TreeMap<>(); + long lastEpoch = Long.MIN_VALUE; + for (WALObject object : objects) { + if (object.epoch() == lastEpoch) { + continue; + } + newIndexMap.put(object.startOffset(), object.epoch()); + lastEpoch = object.epoch(); + indexLargestEpoch = lastEpoch; + } + this.indexMap = newIndexMap; + }, eventLoop); + } + + static class SingleReadTask { + final long epoch; + final long offset; + final int size; + final CompletableFuture cf; + + public SingleReadTask(long epoch, long offset, int size) { + this.epoch = epoch; + this.offset = offset; + this.size = size; + this.cf = new CompletableFuture<>(); + } + } + + static class BatchReadTask { + final DefaultRecordOffset startOffset; + final DefaultRecordOffset endOffset; + final CompletableFuture> cf; + + public BatchReadTask(RecordOffset startOffset, RecordOffset endOffset) { + this.startOffset = DefaultRecordOffset.of(startOffset); + this.endOffset = DefaultRecordOffset.of(endOffset); + this.cf = new CompletableFuture<>(); + } + + @Override + public String toString() { + return "BatchReadTask{" + + "startOffset=" + startOffset + + ", endOffset=" + endOffset + + '}'; + } + } + +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java new file mode 100644 index 0000000000..a3620e5997 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java @@ -0,0 +1,623 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.ByteBufSeqAlloc; +import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.metrics.stats.StorageOperationStats; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.DefaultAppendResult; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.RecoverResult; +import com.automq.stream.s3.wal.ReservationService; +import com.automq.stream.s3.wal.common.RecordHeader; +import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.s3.wal.exception.RuntimeIOException; +import com.automq.stream.s3.wal.exception.WALFencedException; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; +import com.automq.stream.s3.wal.util.WALUtil; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Systems; +import com.automq.stream.utils.Threads; +import com.automq.stream.utils.Time; +import com.automq.stream.utils.threads.EventLoop; +import com.google.common.annotations.VisibleForTesting; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Objects; +import java.util.Queue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentNavigableMap; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.stream.Collectors; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.buffer.Unpooled; + +import static com.automq.stream.s3.ByteBufAlloc.S3_WAL; +import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; +import static com.automq.stream.s3.wal.impl.object.ObjectUtils.DATA_FILE_ALIGN_SIZE; +import static com.automq.stream.s3.wal.impl.object.ObjectUtils.OBJECT_PATH_OFFSET_DELIMITER; +import static com.automq.stream.s3.wal.impl.object.ObjectUtils.ceilAlignOffset; + +public class DefaultWriter implements Writer { + private static final Logger LOGGER = LoggerFactory.getLogger(DefaultWriter.class); + + private static final long DEFAULT_LOCK_WARNING_TIMEOUT = TimeUnit.MILLISECONDS.toNanos(5); + private static final long DEFAULT_UPLOAD_WARNING_TIMEOUT = TimeUnit.SECONDS.toNanos(5); + private static final String OBJECT_PATH_FORMAT = "%s%d" + OBJECT_PATH_OFFSET_DELIMITER + "%d"; // {objectPrefix}/{startOffset}-{endOffset} + private static final ByteBufSeqAlloc BYTE_BUF_ALLOC = new ByteBufSeqAlloc(S3_WAL, 8); + private static final ExecutorService UPLOAD_EXECUTOR = Threads.newFixedThreadPoolWithMonitor(Systems.CPU_CORES, "S3_WAL_UPLOAD", true, LOGGER); + private static final ScheduledExecutorService SCHEDULE = Threads.newSingleThreadScheduledExecutor("S3_WAL_SCHEDULE", true, LOGGER); + + protected final ObjectWALConfig config; + protected final Time time; + protected final ObjectStorage objectStorage; + protected final ReservationService reservationService; + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private List previousObjects = new ArrayList<>(); + private final ConcurrentNavigableMap lastRecordOffset2object = new ConcurrentSkipListMap<>(); + private final String nodePrefix; + private final String objectPrefix; + + private final AtomicLong objectDataBytes = new AtomicLong(); + private final AtomicLong bufferedDataBytes = new AtomicLong(); + + protected volatile boolean closed = true; + protected volatile boolean fenced; + + private Bulk activeBulk = null; + private Bulk lastInActiveBulk = null; + private long lastBulkForceUploadNanos; + private final long batchNanos; + private final long minBulkUploadIntervalNanos; + + private final Queue waitingUploadBulks = new ConcurrentLinkedQueue<>(); + private final Queue uploadingBulks = new ConcurrentLinkedQueue<>(); + + private CompletableFuture callbackCf = CompletableFuture.completedFuture(null); + private final EventLoop callbackExecutor = new EventLoop("S3_WAL_CALLBACK"); + + private final AtomicLong nextOffset = new AtomicLong(); + private final AtomicLong flushedOffset = new AtomicLong(); + private final AtomicLong trimOffset = new AtomicLong(-1); + private CompletableFuture lastTrimCf = CompletableFuture.completedFuture(null); + + public DefaultWriter(Time time, ObjectStorage objectStorage, ObjectWALConfig config) { + this.time = time; + this.objectStorage = objectStorage; + this.reservationService = config.reservationService(); + this.config = config; + this.nodePrefix = ObjectUtils.nodePrefix(config.clusterId(), config.nodeId(), config.type()); + this.objectPrefix = nodePrefix + config.epoch() + "/wal/"; + this.batchNanos = TimeUnit.MILLISECONDS.toNanos(config.batchInterval()); + this.minBulkUploadIntervalNanos = Math.min(TimeUnit.MILLISECONDS.toNanos(10), batchNanos); + this.lastBulkForceUploadNanos = time.nanoseconds(); + if (!(config.openMode() == OpenMode.READ_WRITE || config.openMode() == OpenMode.FAILOVER)) { + throw new IllegalArgumentException("The open mode must be READ_WRITE or FAILOVER, but got " + config.openMode()); + } + } + + public void start() { + // Verify the permission. + reservationService.verify(config.nodeId(), config.epoch(), config.openMode() == OpenMode.FAILOVER) + .thenAccept(result -> { + if (!result) { + fenced = true; + WALFencedException exception = new WALFencedException("Failed to verify the permission with node id: " + config.nodeId() + ", node epoch: " + config.epoch() + ", failover flag: " + config.openMode()); + throw new CompletionException(exception); + } + }) + .join(); + List objects = objectStorage.list(nodePrefix).thenApply(ObjectUtils::parse).join(); + List overlapObjects = ObjectUtils.skipOverlapObjects(objects); + if (!overlapObjects.isEmpty()) { + objectStorage + .delete(overlapObjects.stream() + .map(o -> new ObjectStorage.ObjectPath(o.bucketId(), o.path())) + .collect(Collectors.toList()) + ) + .thenAccept(nil -> LOGGER.info("Delete overlap objects: {}", overlapObjects)); + } + long largestEpoch; + if (!objects.isEmpty() && (largestEpoch = objects.get(objects.size() - 1).epoch()) > config.epoch()) { + LOGGER.warn("Detect newer epoch={} WAL started, exit current WAL start", largestEpoch); + fenced = true; + return; + } + objects.forEach(object -> objectDataBytes.addAndGet(object.length())); + + previousObjects.addAll(objects); + + // TODO: force align before accept new data. + flushedOffset.set(objects.isEmpty() ? 0 : objects.get(objects.size() - 1).endOffset()); + nextOffset.set(flushedOffset.get()); + + startMonitor(); + + closed = false; + } + + @Override + public void close() { + closed = true; + uploadActiveBulk(); + if (lastInActiveBulk != null) { + try { + lastInActiveBulk.completeCf.get(); + } catch (Throwable ex) { + LOGGER.error("Failed to flush records when close.", ex); + } + } + + LOGGER.info("S3WAL Writer is closed."); + } + + List objectList() throws WALFencedException { + checkStatus(); + List list = new ArrayList<>(lastRecordOffset2object.size() + previousObjects.size()); + list.addAll(previousObjects); + list.addAll(lastRecordOffset2object.values()); + return list; + } + + protected void checkStatus() throws WALFencedException { + if (closed) { + throw new IllegalStateException("WAL is closed."); + } + + if (fenced) { + throw new WALFencedException("WAL is fenced."); + } + } + + protected void checkWriteStatus() throws WALFencedException { + checkStatus(); + } + + public CompletableFuture append(StreamRecordBatch streamRecordBatch) throws OverCapacityException { + try { + return append0(streamRecordBatch); + } catch (Throwable ex) { + streamRecordBatch.release(); + if (ex instanceof OverCapacityException) { + throw (OverCapacityException) ex; + } else { + return CompletableFuture.failedFuture(ex); + } + } + } + + @Override + public RecordOffset confirmOffset() { + return DefaultRecordOffset.of(config.epoch(), flushedOffset.get(), 0); + } + + public CompletableFuture append0( + StreamRecordBatch streamRecordBatch) throws OverCapacityException, WALFencedException { + checkWriteStatus(); + + if (bufferedDataBytes.get() > config.maxUnflushedBytes()) { + throw new OverCapacityException(String.format("Max unflushed bytes exceeded %s > %s.", bufferedDataBytes.get(), config.maxUnflushedBytes())); + } + + int dataSize = streamRecordBatch.encoded().readableBytes() + RecordHeader.RECORD_HEADER_SIZE; + if (dataSize > DATA_FILE_ALIGN_SIZE) { + throw new IllegalStateException("Data size exceeded " + dataSize + " > " + DATA_FILE_ALIGN_SIZE); + } + + Record record = new Record(streamRecordBatch, new CompletableFuture<>()); + lock.writeLock().lock(); + try { + if (activeBulk == null) { + activeBulk = new Bulk(nextOffset.get()); + } + if (dataSize + activeBulk.size > DATA_FILE_ALIGN_SIZE) { + uploadActiveBulk(); + this.activeBulk = new Bulk(nextOffset.get()); + } + bufferedDataBytes.addAndGet(dataSize); + activeBulk.add(record); + if (activeBulk.size > config.maxBytesInBatch()) { + uploadActiveBulk(); + } + } finally { + lock.writeLock().unlock(); + } + return record.future.whenComplete((v, throwable) -> { + bufferedDataBytes.addAndGet(-dataSize); + if (throwable != null) { + LOGGER.error("Failed to append record to S3 WAL", throwable); + } + }); + } + + private void forceUploadBulk(Bulk forceBulk) { + lock.writeLock().lock(); + try { + if (forceBulk == this.activeBulk) { + uploadActiveBulk(); + } + } finally { + lock.writeLock().unlock(); + } + } + + private void uploadActiveBulk() { + lock.writeLock().lock(); + try { + if (activeBulk == null) { + return; + } + waitingUploadBulks.add(activeBulk); + nextOffset.set(ObjectUtils.ceilAlignOffset(nextOffset.get() + activeBulk.size)); + lastInActiveBulk = activeBulk; + activeBulk = null; + } finally { + lock.writeLock().unlock(); + } + tryUploadBulkInWaiting(); + } + + @VisibleForTesting + CompletableFuture flush() { + uploadActiveBulk(); + lock.writeLock().lock(); + try { + return lastInActiveBulk == null ? CompletableFuture.completedFuture(null) : lastInActiveBulk.completeCf; + } finally { + lock.writeLock().unlock(); + } + } + + private void tryUploadBulkInWaiting() { + lock.writeLock().lock(); + try { + while (uploadingBulks.size() < config.maxInflightUploadCount()) { + Bulk bulk = waitingUploadBulks.poll(); + if (bulk == null) { + return; + } + uploadingBulks.add(bulk); + UPLOAD_EXECUTOR.submit(() -> uploadBulk0(bulk)); + } + } finally { + lock.writeLock().unlock(); + } + } + + private void uploadBulk0(Bulk bulk) { + try { + long startTime = time.nanoseconds(); + List records = bulk.records; + // Order by + records.sort((o1, o2) -> { + StreamRecordBatch s1 = o1.streamRecordBatch; + StreamRecordBatch s2 = o2.streamRecordBatch; + int rst = Long.compare(s1.getStreamId(), s2.getStreamId()); + if (rst != 0) { + return rst; + } + rst = Long.compare(s1.getBaseOffset(), s2.getBaseOffset()); + return rst; + }); + + long firstOffset = bulk.baseOffset; + long nextOffset = firstOffset; + long lastRecordOffset = nextOffset; + CompositeByteBuf dataBuffer = ByteBufAlloc.compositeByteBuffer(); + for (Record record : records) { + record.offset = nextOffset; + lastRecordOffset = record.offset; + ByteBuf data = record.streamRecordBatch.encoded(); + ByteBuf header = BYTE_BUF_ALLOC.byteBuffer(RECORD_HEADER_SIZE); + header = WALUtil.generateHeader(data, header, 0, nextOffset); + nextOffset += record.size; + dataBuffer.addComponent(true, header); + dataBuffer.addComponent(true, data); + } + + // Build object buffer. + long dataLength = dataBuffer.readableBytes(); + nextOffset = ObjectUtils.ceilAlignOffset(nextOffset); + long endOffset = nextOffset; + + CompositeByteBuf objectBuffer = ByteBufAlloc.compositeByteBuffer(); + WALObjectHeader header = new WALObjectHeader(firstOffset, dataLength, 0, config.nodeId(), config.epoch(), trimOffset.get()); + objectBuffer.addComponent(true, header.marshal()); + objectBuffer.addComponent(true, dataBuffer); + + // Trigger upload. + int objectLength = objectBuffer.readableBytes(); + + // Enable fast retry. + ObjectStorage.WriteOptions writeOptions = new ObjectStorage.WriteOptions().enableFastRetry(true); + String path = String.format(OBJECT_PATH_FORMAT, objectPrefix, firstOffset, endOffset); + FutureUtil.propagate(objectStorage.write(writeOptions, path, objectBuffer), bulk.uploadCf); + long finalLastRecordOffset = lastRecordOffset; + bulk.uploadCf.whenCompleteAsync((rst, ex) -> { + if (ex != null) { + fenced = true; + LOGGER.error("S3WAL upload {} fail", path, ex); + } else { + StorageOperationStats.getInstance().appendWALWriteStats.record(time.nanoseconds() - startTime); + lastRecordOffset2object.put(finalLastRecordOffset, new WALObject(rst.bucket(), path, config.epoch(), firstOffset, endOffset, objectLength)); + objectDataBytes.addAndGet(objectLength); + } + callback(); + }, callbackExecutor); + } catch (Throwable ex) { + bulk.uploadCf.completeExceptionally(ex); + } + } + + private void callback() { + callbackCf = callbackCf.thenComposeAsync(nil -> { + List completedBulks = new ArrayList<>(); + while (true) { + Bulk bulk = uploadingBulks.peek(); + if (bulk == null || !bulk.uploadCf.isDone()) { + break; + } + uploadingBulks.poll(); + completedBulks.add(bulk); + } + if (completedBulks.isEmpty()) { + return CompletableFuture.completedFuture(null); + } + // The inflight uploading bulks count was decreased, then trigger the upload of Bulk in waitingUploadBulks + tryUploadBulkInWaiting(); + return reservationService.verify(config.nodeId(), config.epoch(), config.openMode() == OpenMode.FAILOVER) + .whenComplete((rst, ex) -> { + if (ex != null) { + LOGGER.error("Unexpected S3WAL lease check fail. Make the WAL fenced", ex); + fenced = true; + } else if (!rst) { + LOGGER.warn("The S3WAL is fenced by another nodes. Fail the following append request"); + fenced = true; + } + }).exceptionally(ex -> false) + .thenAcceptAsync(rst -> { + if (fenced) { + Throwable ex = new WALFencedException(); + for (Bulk bulk : completedBulks) { + bulk.complete(ex); + } + } else { + for (Bulk bulk : completedBulks) { + flushedOffset.set(ObjectUtils.ceilAlignOffset(bulk.endOffset())); + bulk.complete(null); + } + } + }, callbackExecutor); + }, callbackExecutor).exceptionally(ex -> { + LOGGER.error("Unexpected S3WAL callback fail", ex); + return null; + }); + } + + public CompletableFuture reset() throws WALFencedException { + long nextOffset = this.nextOffset.get(); + if (nextOffset == 0) { + return CompletableFuture.completedFuture(null); + } + // The next offset is the next record's offset. + return trim0(nextOffset - 1); + } + + // Trim objects where the last offset is less than or equal to the given offset. + public CompletableFuture trim(RecordOffset recordOffset) throws WALFencedException { + long newStartOffset = ((DefaultRecordOffset) recordOffset).offset(); + return trim0(newStartOffset); + } + + @Override + public Iterator recover() { + try { + return new RecoverIterator(objectList(), objectStorage, config.readAheadObjectCount()); + } catch (WALFencedException e) { + LOGGER.error("Recover S3 WAL failed, due to unrecoverable exception.", e); + return new Iterator<>() { + @Override + public boolean hasNext() { + throw new RuntimeIOException(e); + } + + @Override + public RecoverResult next() { + throw new RuntimeIOException(e); + } + }; + } + } + + // Trim objects where the last offset is less than or equal to the given offset. + public CompletableFuture trim0(long inclusiveTrimRecordOffset) throws WALFencedException { + checkStatus(); + List deleteObjectList = new ArrayList<>(); + AtomicLong deletedObjectSize = new AtomicLong(); + CompletableFuture persistTrimOffsetCf; + lock.writeLock().lock(); + try { + if (trimOffset.get() >= inclusiveTrimRecordOffset) { + return lastTrimCf; + } + trimOffset.set(inclusiveTrimRecordOffset); + // We cannot force upload an empty wal object cause of the recover workflow don't accept an empty wal object. + // So we use a fake record to trigger the wal object upload. + persistTrimOffsetCf = append(StreamRecordBatch.of(-1L, -1L, 0, 0, Unpooled.EMPTY_BUFFER)); + lastTrimCf = persistTrimOffsetCf.thenCompose(nil -> { + Long lastFlushedRecordOffset = lastRecordOffset2object.isEmpty() ? null : lastRecordOffset2object.lastKey(); + if (lastFlushedRecordOffset != null) { + lastRecordOffset2object.headMap(inclusiveTrimRecordOffset, true) + .forEach((lastRecordOffset, object) -> { + if (Objects.equals(lastRecordOffset, lastFlushedRecordOffset)) { + // skip the last object to prevent wal offset reset back to zero + // when there is no object could be used to calculate the nextOffset after restart. + return; + } + deleteObjectList.add(new ObjectStorage.ObjectPath(object.bucketId(), object.path())); + deletedObjectSize.addAndGet(object.length()); + lastRecordOffset2object.remove(lastRecordOffset); + }); + } + + if (!previousObjects.isEmpty()) { + boolean skipTheLastObject = deleteObjectList.isEmpty(); + List list = new ArrayList<>(previousObjects.size()); + for (int i = 0; i < previousObjects.size() - (skipTheLastObject ? 1 : 0); i++) { + WALObject object = previousObjects.get(i); + if (object.endOffset() > inclusiveTrimRecordOffset) { + break; + } + list.add(new ObjectStorage.ObjectPath(object.bucketId(), object.path())); + deletedObjectSize.addAndGet(object.length()); + } + previousObjects = new ArrayList<>(previousObjects.subList(list.size(), previousObjects.size())); + deleteObjectList.addAll(list); + } + if (deleteObjectList.isEmpty()) { + return CompletableFuture.completedFuture(null); + } + + return objectStorage.delete(deleteObjectList).whenComplete((v, throwable) -> { + objectDataBytes.addAndGet(-1 * deletedObjectSize.get()); + // Never fail the delete task, the under layer storage will retry forever. + if (throwable != null) { + LOGGER.error("Failed to delete objects when trim S3 WAL: {}", deleteObjectList, throwable); + } + SCHEDULE.schedule(() -> { + // - Try to Delete the objects again after 30 seconds to avoid object leak because of underlying fast retry + objectStorage.delete(deleteObjectList); + }, 10, TimeUnit.SECONDS); + }); + }); + return lastTrimCf; + } catch (Throwable e) { + return CompletableFuture.failedFuture(e); + } finally { + lock.writeLock().unlock(); + } + } + + private void startMonitor() { + SCHEDULE.scheduleWithFixedDelay(() -> { + try { + long count = uploadingBulks.stream() + .filter(bulk -> time.nanoseconds() - bulk.startNanos > DEFAULT_UPLOAD_WARNING_TIMEOUT) + .count(); + if (count > 0) { + LOGGER.error("Found {} pending upload tasks exceed 5s.", count); + } + } catch (Throwable ignore) { + } + }, 1, 1, TimeUnit.SECONDS); + } + + class Bulk { + private int size; + private final long baseOffset; + private final List records = new ArrayList<>(1024); + private final long startNanos; + final CompletableFuture uploadCf = new CompletableFuture<>(); + final CompletableFuture completeCf = new CompletableFuture<>(); + + public Bulk(long baseOffset) { + this.startNanos = time.nanoseconds(); + this.baseOffset = baseOffset; + long forceUploadDelayNanos = Math.min( + Math.max( + // Try batch the requests in a short time window to save the PUT API. + minBulkUploadIntervalNanos, + lastBulkForceUploadNanos + batchNanos - startNanos + ), + batchNanos); + lastBulkForceUploadNanos = startNanos + forceUploadDelayNanos; + SCHEDULE.schedule(() -> forceUploadBulk(this), forceUploadDelayNanos, TimeUnit.NANOSECONDS); + } + + public void add(Record record) { + records.add(record); + size += record.size; + } + + public int size() { + return size; + } + + public long baseOffset() { + return baseOffset; + } + + public long endOffset() { + return baseOffset + size; + } + + public void complete(Throwable ex) { + if (ex != null) { + records.forEach(record -> record.future.completeExceptionally(ex)); + } else { + for (int idx = 0; idx < records.size(); idx++) { + Record record = records.get(idx); + // Requests for data starting at or beyond objectSize will return a 416 error. + // So we jump the last record's nextOffset to the ceil-aligned offset. + long nextOffset = (idx == records.size() - 1) ? ceilAlignOffset(record.offset + record.size) : record.offset + record.size; + record.future.complete( + new DefaultAppendResult( + DefaultRecordOffset.of(config.epoch(), record.offset, record.size), + DefaultRecordOffset.of(config.epoch(), nextOffset, 0) + ) + ); + } + } + completeCf.complete(null); + } + } + + protected static class Record { + public final StreamRecordBatch streamRecordBatch; + public final CompletableFuture future; + public final int size; + public long offset; + + public Record(StreamRecordBatch streamRecordBatch, CompletableFuture future) { + this.streamRecordBatch = streamRecordBatch; + this.future = future; + this.size = streamRecordBatch.encoded().readableBytes() + RECORD_HEADER_SIZE; + } + } + +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/NoopWriter.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/NoopWriter.java new file mode 100644 index 0000000000..5c82720674 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/NoopWriter.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.RecoverResult; +import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.s3.wal.exception.WALFencedException; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; + +import java.util.Iterator; +import java.util.concurrent.CompletableFuture; + +public class NoopWriter implements Writer { + + @Override + public void start() { + // No-op + } + + @Override + public void close() { + // No-op + } + + @Override + public CompletableFuture append(StreamRecordBatch streamRecordBatch) throws OverCapacityException { + return CompletableFuture.failedFuture(new UnsupportedOperationException("append is not supported")); + } + + @Override + public RecordOffset confirmOffset() { + return DefaultRecordOffset.of(0, 0, 0); + } + + @Override + public CompletableFuture reset() throws WALFencedException { + return CompletableFuture.failedFuture(new UnsupportedOperationException("reset is not supported")); + } + + @Override + public CompletableFuture trim(RecordOffset recordOffset) throws WALFencedException { + return CompletableFuture.failedFuture(new UnsupportedOperationException("trim is not supported")); + } + + @Override + public Iterator recover() { + throw new UnsupportedOperationException("recover is not supported"); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java new file mode 100644 index 0000000000..94dc0f8979 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java @@ -0,0 +1,108 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.Constants; +import com.automq.stream.s3.network.ThrottleStrategy; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.wal.ReservationService; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class ObjectReservationService implements ReservationService { + private static final Logger LOGGER = LoggerFactory.getLogger(ObjectReservationService.class); + + public static final int S3_RESERVATION_OBJECT_MAGIC_CODE = 0x12345678; + public static final int S3_RESERVATION_OBJECT_LENGTH = 4 // magic code + + 8 // node id + + 8 // node epoch + + 1; // failover flag + + private final String clusterId; + private final ObjectStorage objectStorage; + private final short bucketId; + + private final ConcurrentMap nodeIdPathMap; + + public ObjectReservationService(String clusterId, ObjectStorage objectStorage, short bucketId) { + this.clusterId = clusterId; + this.objectStorage = objectStorage; + this.nodeIdPathMap = new ConcurrentHashMap<>(); + this.bucketId = bucketId; + } + + private String path(long nodeId) { + return nodeIdPathMap.computeIfAbsent(nodeId, node -> "reservation/" + Constants.DEFAULT_NAMESPACE + clusterId + "/" + node); + } + + // Visible for testing + protected CompletableFuture verify(long nodeId, ByteBuf target) { + ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions().throttleStrategy(ThrottleStrategy.BYPASS).bucket(bucketId); + return objectStorage.rangeRead(options, path(nodeId), 0, S3_RESERVATION_OBJECT_LENGTH) + .thenApply(bytes -> { + try { + ByteBuf slice = bytes.slice(); + slice.readInt(); + if (bytes.readableBytes() != S3_RESERVATION_OBJECT_LENGTH) { + return false; + } + return bytes.equals(target); + } finally { + bytes.release(); + target.release(); + } + }) + .exceptionally(e -> { + LOGGER.error("Check reservation object failed:", e); + return false; + }); + } + + @Override + public CompletableFuture verify(long nodeId, long epoch, boolean failover) { + ByteBuf target = Unpooled.buffer(S3_RESERVATION_OBJECT_LENGTH); + target.writeInt(S3_RESERVATION_OBJECT_MAGIC_CODE); + target.writeLong(nodeId); + target.writeLong(epoch); + target.writeBoolean(failover); + return verify(nodeId, target); + } + + @Override + public CompletableFuture acquire(long nodeId, long epoch, boolean failover) { + LOGGER.info("Acquire permission for node: {}, epoch: {}, failover: {}", nodeId, epoch, failover); + ByteBuf target = ByteBufAlloc.byteBuffer(S3_RESERVATION_OBJECT_LENGTH); + target.writeInt(S3_RESERVATION_OBJECT_MAGIC_CODE); + target.writeLong(nodeId); + target.writeLong(epoch); + target.writeBoolean(failover); + ObjectStorage.WriteOptions options = new ObjectStorage.WriteOptions().throttleStrategy(ThrottleStrategy.BYPASS); + return objectStorage.write(options, path(nodeId), target).thenApply(rst -> null); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectUtils.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectUtils.java new file mode 100644 index 0000000000..581082e833 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectUtils.java @@ -0,0 +1,142 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.Constants; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.wal.common.RecordHeader; + +import org.apache.commons.codec.digest.DigestUtils; +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Locale; + +import io.netty.buffer.ByteBuf; + +import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; + +public class ObjectUtils { + public static final long DATA_FILE_ALIGN_SIZE = 64L * 1024 * 1024; // 64MiB + private static final Logger LOGGER = LoggerFactory.getLogger(ObjectUtils.class); + public static final String OBJECT_PATH_OFFSET_DELIMITER = "-"; + + public static List parse(List objects) { + List walObjects = new ArrayList<>(objects.size()); + for (ObjectStorage.ObjectInfo object : objects) { + // v0: md5hex(nodeId)/clusterId/nodeId/epoch/wal/startOffset + // v1: md5hex(nodeId)/clusterId/nodeId/epoch/wal/startOffset_endOffset + String path = object.key(); + String[] parts = path.split("/"); + try { + long epoch = Long.parseLong(parts[parts.length - 3]); + long length = object.size(); + String rawOffset = parts[parts.length - 1]; + + WALObject walObject; + if (rawOffset.contains(OBJECT_PATH_OFFSET_DELIMITER)) { + // v1 format: {startOffset}-{endOffset} + long startOffset = Long.parseLong(rawOffset.substring(0, rawOffset.indexOf(OBJECT_PATH_OFFSET_DELIMITER))); + long endOffset = Long.parseLong(rawOffset.substring(rawOffset.indexOf(OBJECT_PATH_OFFSET_DELIMITER) + 1)); + walObject = new WALObject(object.bucketId(), path, epoch, startOffset, endOffset, length); + } else { + // v0 format: {startOffset} + long startOffset = Long.parseLong(rawOffset); + walObject = new WALObject(object.bucketId(), path, epoch, startOffset, length); + } + walObjects.add(walObject); + } catch (NumberFormatException e) { + // Ignore invalid path + LOGGER.warn("Invalid WAL object: {}", path); + } + } + walObjects.sort(Comparator.comparingLong(WALObject::epoch).thenComparingLong(WALObject::startOffset)); + return walObjects; + } + + /** + * Remove overlap objects. + * + * @return overlap objects. + */ + public static List skipOverlapObjects(List objects) { + List overlapObjects = new ArrayList<>(); + WALObject lastObject = null; + for (WALObject object : objects) { + if (lastObject == null) { + lastObject = object; + continue; + } + if (lastObject.epoch() != object.epoch()) { + if (lastObject.endOffset() > object.startOffset()) { + // maybe the old epoch node write dirty object after it was fenced. + overlapObjects.add(lastObject); + } + } + } + overlapObjects.forEach(objects::remove); + return overlapObjects; + } + + public static String nodePrefix(String clusterId, int nodeId, String type) { + return DigestUtils.md5Hex(String.valueOf(nodeId)).toUpperCase(Locale.ROOT) + "/" + Constants.DEFAULT_NAMESPACE + clusterId + "/" + nodeId + (StringUtils.isBlank(type) ? "" : ("_" + type)) + "/"; + } + + public static String nodePrefix(String clusterId, int nodeId) { + return nodePrefix(clusterId, nodeId, null); + } + + public static String genObjectPathV0(String nodePrefix, long epoch, long objectStartOffset) { + return nodePrefix + epoch + "/wal/" + objectStartOffset; + } + + public static String genObjectPathV1(String nodePrefix, long epoch, long objectStartOffset, long objectEndOffset) { + return nodePrefix + epoch + "/wal/" + objectStartOffset + OBJECT_PATH_OFFSET_DELIMITER + objectEndOffset; + } + + public static String genObjectPathV1(String nodePrefix, long epoch, long objectStartOffset) { + long endOffset = objectStartOffset + DATA_FILE_ALIGN_SIZE; + return nodePrefix + epoch + "/wal/" + objectStartOffset + OBJECT_PATH_OFFSET_DELIMITER + endOffset; + } + + public static long floorAlignOffset(long offset) { + return offset / DATA_FILE_ALIGN_SIZE * DATA_FILE_ALIGN_SIZE; + } + + public static long ceilAlignOffset(long offset) { + return (offset + DATA_FILE_ALIGN_SIZE - 1) / DATA_FILE_ALIGN_SIZE * DATA_FILE_ALIGN_SIZE; + } + + public static StreamRecordBatch decodeRecordBuf(ByteBuf dataBuffer) { + // TODO: thread local to avoid alloc + ByteBuf recordHeaderBuf = dataBuffer.readBytes(RECORD_HEADER_SIZE); + RecordHeader header = new RecordHeader(recordHeaderBuf); + recordHeaderBuf.release(); + if (header.getMagicCode() != RecordHeader.RECORD_HEADER_DATA_MAGIC_CODE) { + throw new IllegalStateException("Invalid magic code in record header."); + } + return StreamRecordBatch.parse(dataBuffer, false); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java index e52c416d89..00d54d1d3f 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java @@ -1,21 +1,33 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.impl.object; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.s3.wal.ReservationService; import com.automq.stream.utils.IdURI; import org.apache.commons.lang3.StringUtils; public class ObjectWALConfig { + private final String uri; + private final ReservationService reservationService; private final long batchInterval; private final long maxBytesInBatch; private final long maxUnflushedBytes; @@ -24,17 +36,20 @@ public class ObjectWALConfig { private final String clusterId; private final int nodeId; private final long epoch; - private final boolean failover; + private final OpenMode openMode; private final short bucketId; - private final boolean strictBatchLimit; + private final String type; public static Builder builder() { return new Builder(); } - public ObjectWALConfig(long batchInterval, long maxBytesInBatch, long maxUnflushedBytes, int maxInflightUploadCount, - int readAheadObjectCount, String clusterId, int nodeId, long epoch, boolean failover, short bucketId, - boolean strictBatchLimit) { + public ObjectWALConfig(String uri, ReservationService reservationService, long batchInterval, long maxBytesInBatch, + long maxUnflushedBytes, int maxInflightUploadCount, + int readAheadObjectCount, String clusterId, int nodeId, long epoch, OpenMode openMode, short bucketId, + String type) { + this.uri = uri; + this.reservationService = reservationService; this.batchInterval = batchInterval; this.maxBytesInBatch = maxBytesInBatch; this.maxUnflushedBytes = maxUnflushedBytes; @@ -43,9 +58,17 @@ public ObjectWALConfig(long batchInterval, long maxBytesInBatch, long maxUnflush this.clusterId = clusterId; this.nodeId = nodeId; this.epoch = epoch; - this.failover = failover; + this.openMode = openMode; this.bucketId = bucketId; - this.strictBatchLimit = strictBatchLimit; + this.type = type; + } + + public String uri() { + return uri; + } + + public ReservationService reservationService() { + return this.reservationService; } public long batchInterval() { @@ -80,35 +103,55 @@ public long epoch() { return epoch; } - public boolean failover() { - return failover; + public OpenMode openMode() { + return openMode; } public short bucketId() { return bucketId; } - public boolean strictBatchLimit() { - return strictBatchLimit; + public String type() { + return type; + } + + @Override + public String toString() { + return "ObjectWALConfig{" + + "batchInterval=" + batchInterval + + ", maxBytesInBatch=" + maxBytesInBatch + + ", maxUnflushedBytes=" + maxUnflushedBytes + + ", maxInflightUploadCount=" + maxInflightUploadCount + + ", readAheadObjectCount=" + readAheadObjectCount + + ", clusterId='" + clusterId + '\'' + + ", nodeId=" + nodeId + + ", epoch=" + epoch + + ", openMode=" + openMode + + ", bucketId=" + bucketId + + ", type='" + type + '\'' + + '}'; } public static final class Builder { - private long batchInterval = 100; // 100ms - private long maxBytesInBatch = 4 * 1024 * 1024L; // 4MB + private String uri = ""; + private ReservationService reservationService = ReservationService.NOOP; + private long batchInterval = 250; // 250ms + private long maxBytesInBatch = 8 * 1024 * 1024L; // 8MB private long maxUnflushedBytes = 1024 * 1024 * 1024L; // 1GB private int maxInflightUploadCount = 50; private int readAheadObjectCount = 4; private String clusterId; private int nodeId; private long epoch; - private boolean failover; + private OpenMode openMode = OpenMode.READ_WRITE; private short bucketId; - private boolean strictBatchLimit = false; + private String type = ""; private Builder() { } public Builder withURI(IdURI uri) { + this.uri = uri.toString(); withBucketId(uri.id()); String batchInterval = uri.extensionString("batchInterval"); @@ -131,10 +174,15 @@ public Builder withURI(IdURI uri) { if (StringUtils.isNumeric(readAheadObjectCount)) { withReadAheadObjectCount(Integer.parseInt(readAheadObjectCount)); } - String strictBatchLimit = uri.extensionString("strictBatchLimit"); - if (StringUtils.isNumeric(strictBatchLimit)) { - withStrictBatchLimit(Boolean.parseBoolean(strictBatchLimit)); - } + return this; + } + + public Builder withURI(String uri) { + return withURI(IdURI.parse(uri)); + } + + public Builder withReservationService(ReservationService reservationService) { + this.reservationService = reservationService; return this; } @@ -186,8 +234,8 @@ public Builder withEpoch(long epoch) { return this; } - public Builder withFailover(boolean failover) { - this.failover = failover; + public Builder withOpenMode(OpenMode openMode) { + this.openMode = openMode; return this; } @@ -196,13 +244,13 @@ public Builder withBucketId(short bucketId) { return this; } - public Builder withStrictBatchLimit(boolean strictBatchLimit) { - this.strictBatchLimit = strictBatchLimit; + public Builder withType(String type) { + this.type = type; return this; } public ObjectWALConfig build() { - return new ObjectWALConfig(batchInterval, maxBytesInBatch, maxUnflushedBytes, maxInflightUploadCount, readAheadObjectCount, clusterId, nodeId, epoch, failover, bucketId, strictBatchLimit); + return new ObjectWALConfig(uri, reservationService, batchInterval, maxBytesInBatch, maxUnflushedBytes, maxInflightUploadCount, readAheadObjectCount, clusterId, nodeId, epoch, openMode, bucketId, type); } } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java index e8001b7e18..a6ce20d24d 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALService.java @@ -1,87 +1,75 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.impl.object; -import com.automq.stream.ByteBufSeqAlloc; -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.network.ThrottleStrategy; +import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.trace.context.TraceContext; import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.s3.wal.RecordOffset; import com.automq.stream.s3.wal.RecoverResult; import com.automq.stream.s3.wal.WriteAheadLog; -import com.automq.stream.s3.wal.common.AppendResultImpl; -import com.automq.stream.s3.wal.common.Record; -import com.automq.stream.s3.wal.common.RecordHeader; -import com.automq.stream.s3.wal.common.RecoverResultImpl; import com.automq.stream.s3.wal.common.WALMetadata; import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.RuntimeIOException; -import com.automq.stream.s3.wal.exception.WALFencedException; -import com.automq.stream.s3.wal.util.WALUtil; import com.automq.stream.utils.Time; -import org.apache.commons.lang3.exception.ExceptionUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; -import java.util.ArrayDeque; import java.util.Iterator; import java.util.List; -import java.util.Objects; -import java.util.Queue; import java.util.concurrent.CompletableFuture; -import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; -import io.netty.buffer.Unpooled; - -import static com.automq.stream.s3.ByteBufAlloc.S3_WAL; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_WITHOUT_CRC_SIZE; - public class ObjectWALService implements WriteAheadLog { private static final Logger log = LoggerFactory.getLogger(ObjectWALService.class); - private static final ByteBufSeqAlloc BYTE_BUF_ALLOC = new ByteBufSeqAlloc(S3_WAL, 8); protected ObjectStorage objectStorage; protected ObjectWALConfig config; - protected RecordAccumulator accumulator; + protected final Writer writer; + protected final DefaultReader reader; public ObjectWALService(Time time, ObjectStorage objectStorage, ObjectWALConfig config) { this.objectStorage = objectStorage; this.config = config; - - this.accumulator = new RecordAccumulator(time, objectStorage, config); - } - - // Visible for testing. - protected RecordAccumulator accumulator() { - return accumulator; + if (config.openMode() == OpenMode.READ_WRITE || config.openMode() == OpenMode.FAILOVER) { + this.writer = new DefaultWriter(time, objectStorage, config); + } else { + this.writer = new NoopWriter(); + } + this.reader = new DefaultReader(objectStorage, config.clusterId(), config.nodeId(), config.type(), time); } @Override public WriteAheadLog start() throws IOException { log.info("Start S3 WAL."); - accumulator.start(); + writer.start(); return this; } @Override public void shutdownGracefully() { log.info("Shutdown S3 WAL."); - accumulator.close(); + writer.close(); } @Override @@ -90,79 +78,41 @@ public WALMetadata metadata() { } @Override - public AppendResult append(TraceContext context, ByteBuf data, int crc) throws OverCapacityException { - ByteBuf header = BYTE_BUF_ALLOC.byteBuffer(RECORD_HEADER_SIZE); - assert header.refCnt() == 1; - - final CompletableFuture appendResultFuture = new CompletableFuture<>(); - try { - final long recordSize = RECORD_HEADER_SIZE + data.readableBytes(); - - long expectedWriteOffset = accumulator.append(recordSize, start -> { - CompositeByteBuf recordByteBuf = ByteBufAlloc.compositeByteBuffer(); - Record record = WALUtil.generateRecord(data, header, 0, start, true); - recordByteBuf.addComponents(true, record.header(), record.body()); - return recordByteBuf; - }, appendResultFuture); + public String uri() { + return config.uri(); + } - return new AppendResultImpl(expectedWriteOffset, appendResultFuture); - } catch (Exception e) { - // Make sure the header buffer and data buffer is released. - if (header.refCnt() > 0) { - header.release(); - } else { - log.error("[Bug] The header buffer is already released.", e); - } + @Override + public CompletableFuture append(TraceContext context, + StreamRecordBatch streamRecordBatch) throws OverCapacityException { + return writer.append(streamRecordBatch); + } - if (data.refCnt() > 0) { - data.release(); - } else { - log.error("[Bug] The data buffer is already released.", e); - } + @Override + public CompletableFuture get(RecordOffset recordOffset) { + return reader.get(recordOffset); + } - // Complete the future with exception, ensure the whenComplete method is executed. - appendResultFuture.completeExceptionally(e); - Throwable cause = ExceptionUtils.getRootCause(e); - if (cause instanceof OverCapacityException) { - if (((OverCapacityException) cause).error()) { - log.warn("Append record to S3 WAL failed, due to accumulator is full.", e); - } else { - log.warn("S3 WAL accumulator is full, try to trigger an upload and trim the WAL", e); - } + @Override + public CompletableFuture> get(RecordOffset startOffset, RecordOffset endOffset) { + return reader.get(startOffset, endOffset); + } - throw new OverCapacityException("Append record to S3 WAL failed, due to accumulator is full: " + cause.getMessage()); - } else { - log.error("Append record to S3 WAL failed, due to unrecoverable exception.", e); - return new AppendResultImpl(-1, appendResultFuture); - } - } + @Override + public RecordOffset confirmOffset() { + return writer.confirmOffset(); } @Override public Iterator recover() { - try { - return new RecoverIterator(accumulator.objectList(), objectStorage, config.readAheadObjectCount()); - } catch (WALFencedException e) { - log.error("Recover S3 WAL failed, due to unrecoverable exception.", e); - return new Iterator<>() { - @Override - public boolean hasNext() { - throw new RuntimeIOException(e); - } - - @Override - public RecoverResult next() { - throw new RuntimeIOException(e); - } - }; - } + return writer.recover(); } @Override public CompletableFuture reset() { log.info("Reset S3 WAL"); try { - return accumulator.reset(); + return writer.reset(); } catch (Throwable e) { log.error("Reset S3 WAL failed, due to unrecoverable exception.", e); return CompletableFuture.failedFuture(e); @@ -170,129 +120,18 @@ public CompletableFuture reset() { } @Override - public CompletableFuture trim(long offset) { + public CompletableFuture trim(RecordOffset offset) { log.info("Trim S3 WAL to offset: {}", offset); try { - return accumulator.trim(offset); + return writer.trim(offset); } catch (Throwable e) { log.error("Trim S3 WAL failed, due to unrecoverable exception.", e); return CompletableFuture.failedFuture(e); } } - public static class RecoverIterator implements Iterator { - private final ObjectStorage objectStorage; - private final int readAheadObjectSize; - - private final List objectList; - private final Queue> readAheadQueue; - - private int nextIndex = 0; - private ByteBuf dataBuffer = Unpooled.EMPTY_BUFFER; - - public RecoverIterator(List objectList, ObjectStorage objectStorage, - int readAheadObjectSize) { - this.objectList = objectList; - this.objectStorage = objectStorage; - this.readAheadObjectSize = readAheadObjectSize; - this.readAheadQueue = new ArrayDeque<>(readAheadObjectSize); - - // Fill the read ahead queue. - for (int i = 0; i < readAheadObjectSize; i++) { - tryReadAhead(); - } - } - - @Override - public boolean hasNext() { - return dataBuffer.isReadable() || !readAheadQueue.isEmpty() || nextIndex < objectList.size(); - } - - private void loadNextBuffer(boolean skipStickyRecord) { - // Please call hasNext() before calling loadNextBuffer(). - byte[] buffer = Objects.requireNonNull(readAheadQueue.poll()).join(); - dataBuffer = Unpooled.wrappedBuffer(buffer); - - // Check header - WALObjectHeader header = WALObjectHeader.unmarshal(dataBuffer); - dataBuffer.skipBytes(WALObjectHeader.WAL_HEADER_SIZE); - - if (skipStickyRecord && header.stickyRecordLength() != 0) { - dataBuffer.skipBytes((int) header.stickyRecordLength()); - } - } - - private void tryReadAhead() { - if (readAheadQueue.size() < readAheadObjectSize && nextIndex < objectList.size()) { - RecordAccumulator.WALObject object = objectList.get(nextIndex++); - ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions().throttleStrategy(ThrottleStrategy.BYPASS).bucket(object.bucketId()); - CompletableFuture readFuture = objectStorage.rangeRead(options, object.path(), 0, object.length()) - .thenApply(buffer -> { - // Copy the result buffer and release it. - byte[] bytes = new byte[buffer.readableBytes()]; - buffer.readBytes(bytes); - buffer.release(); - return bytes; - }); - readAheadQueue.add(readFuture); - } - } - - @Override - public RecoverResult next() { - // If there is no more data to read, return null. - if (!hasNext()) { - return null; - } - - if (!dataBuffer.isReadable()) { - loadNextBuffer(true); - } - - // Try to read next object. - tryReadAhead(); - - ByteBuf recordHeaderBuf = dataBuffer.readBytes(RECORD_HEADER_SIZE); - RecordHeader header = RecordHeader.unmarshal(recordHeaderBuf); - - if (header.getRecordHeaderCRC() != WALUtil.crc32(recordHeaderBuf, RECORD_HEADER_WITHOUT_CRC_SIZE)) { - recordHeaderBuf.release(); - throw new IllegalStateException("Record header crc check failed."); - } - recordHeaderBuf.release(); - - if (header.getMagicCode() != RecordHeader.RECORD_HEADER_MAGIC_CODE) { - throw new IllegalStateException("Invalid magic code in record header."); - } - - int length = header.getRecordBodyLength(); - - ByteBuf recordBuf = ByteBufAlloc.byteBuffer(length); - if (dataBuffer.readableBytes() < length) { - // Read the remain data and release the buffer. - dataBuffer.readBytes(recordBuf, dataBuffer.readableBytes()); - dataBuffer.release(); - - // Read from next buffer. - if (!hasNext()) { - throw new IllegalStateException("[Bug] There is a record part but no more data to read."); - } - loadNextBuffer(false); - dataBuffer.readBytes(recordBuf, length - recordBuf.readableBytes()); - } else { - dataBuffer.readBytes(recordBuf, length); - } - - if (!dataBuffer.isReadable()) { - dataBuffer.release(); - } - - if (header.getRecordBodyCRC() != WALUtil.crc32(recordBuf)) { - recordBuf.release(); - throw new IllegalStateException("Record body crc check failed."); - } - - return new RecoverResultImpl(recordBuf, header.getRecordBodyCRC()); - } + @Override + public String toString() { + return String.format("ObjectWALService{%s@%s-%s-%s}", config.bucketId(), config.nodeId(), config.epoch(), config.type()); } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecordAccumulator.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecordAccumulator.java deleted file mode 100644 index 2e43f29b2b..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecordAccumulator.java +++ /dev/null @@ -1,605 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.object; - -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.Constants; -import com.automq.stream.s3.operator.ObjectStorage; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.common.RecordHeader; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.WALFencedException; -import com.automq.stream.s3.wal.metrics.ObjectWALMetricsManager; -import com.automq.stream.utils.Threads; -import com.automq.stream.utils.Time; - -import org.apache.commons.codec.digest.DigestUtils; -import org.apache.commons.lang3.exception.ExceptionUtils; -import org.apache.commons.lang3.tuple.Pair; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.Closeable; -import java.util.ArrayList; -import java.util.Comparator; -import java.util.LinkedList; -import java.util.List; -import java.util.Locale; -import java.util.PriorityQueue; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentLinkedDeque; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.ConcurrentNavigableMap; -import java.util.concurrent.ConcurrentSkipListMap; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.ReentrantReadWriteLock; -import java.util.function.Function; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; - -public class RecordAccumulator implements Closeable { - private static final Logger log = LoggerFactory.getLogger(RecordAccumulator.class); - - private static final long DEFAULT_LOCK_WARNING_TIMEOUT = TimeUnit.MILLISECONDS.toNanos(5); - private static final long DEFAULT_UPLOAD_WARNING_TIMEOUT = TimeUnit.SECONDS.toNanos(5); - protected final ObjectWALConfig config; - protected final Time time; - protected final ObjectStorage objectStorage; - private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); - private final ConcurrentNavigableMap> uploadMap = new ConcurrentSkipListMap<>(); - private final ConcurrentNavigableMap, WALObject> previousObjectMap = new ConcurrentSkipListMap<>(); - private final ConcurrentNavigableMap objectMap = new ConcurrentSkipListMap<>(); - private final String nodePrefix; - private final String objectPrefix; - private final ScheduledExecutorService executorService; - private final ScheduledExecutorService monitorService; - private final ExecutorService callbackService; - private final ConcurrentMap, Long> pendingFutureMap = new ConcurrentHashMap<>(); - private final AtomicLong objectDataBytes = new AtomicLong(); - private final AtomicLong bufferedDataBytes = new AtomicLong(); - protected volatile boolean closed = true; - protected volatile boolean fenced; - private final ConcurrentLinkedDeque bufferQueue = new ConcurrentLinkedDeque<>(); - private volatile long lastUploadTimestamp = System.currentTimeMillis(); - private final AtomicLong nextOffset = new AtomicLong(); - private final AtomicLong flushedOffset = new AtomicLong(); - - public RecordAccumulator(Time time, ObjectStorage objectStorage, - ObjectWALConfig config) { - this.time = time; - this.objectStorage = objectStorage; - this.config = config; - this.nodePrefix = DigestUtils.md5Hex(String.valueOf(config.nodeId())).toUpperCase(Locale.ROOT) + "/" + Constants.DEFAULT_NAMESPACE + config.clusterId() + "/" + config.nodeId() + "/"; - this.objectPrefix = nodePrefix + config.epoch() + "/wal/"; - this.executorService = Threads.newSingleThreadScheduledExecutor("s3-wal-schedule", true, log); - this.monitorService = Threads.newSingleThreadScheduledExecutor("s3-wal-monitor", true, log); - this.callbackService = Threads.newFixedThreadPoolWithMonitor(4, "s3-wal-callback", false, log); - - ObjectWALMetricsManager.setInflightUploadCountSupplier(() -> (long) pendingFutureMap.size()); - ObjectWALMetricsManager.setBufferedDataInBytesSupplier(bufferedDataBytes::get); - ObjectWALMetricsManager.setObjectDataInBytesSupplier(objectDataBytes::get); - } - - public void start() { - objectStorage.list(nodePrefix) - .thenAccept(objectList -> objectList.forEach(object -> { - String path = object.key(); - String[] parts = path.split("/"); - try { - long firstOffset = Long.parseLong(parts[parts.length - 1]); - long epoch = Long.parseLong(parts[parts.length - 3]); - - // Skip the object if it belongs to a later epoch. - if (epoch > config.epoch()) { - return; - } - - long length = object.size(); - long endOffset = firstOffset + length - WALObjectHeader.WAL_HEADER_SIZE; - - if (epoch != config.epoch()) { - previousObjectMap.put(Pair.of(epoch, endOffset), new WALObject(object.bucketId(), path, firstOffset, length)); - } else { - objectMap.put(endOffset, new WALObject(object.bucketId(), path, firstOffset, length)); - } - objectDataBytes.addAndGet(length); - } catch (NumberFormatException e) { - // Ignore invalid path - log.warn("Found invalid wal object: {}", path); - } - })) - .join(); - - flushedOffset.set(objectMap.isEmpty() ? 0 : objectMap.lastKey()); - nextOffset.set(flushedOffset.get()); - - // Trigger upload periodically. - executorService.scheduleWithFixedDelay(() -> { - long startTime = time.nanoseconds(); - if (fenced - || bufferQueue.isEmpty() - || System.currentTimeMillis() - lastUploadTimestamp < config.batchInterval()) { - return; - } - - lock.writeLock().lock(); - try { - if (time.nanoseconds() - startTime > DEFAULT_LOCK_WARNING_TIMEOUT) { - log.warn("Failed to acquire lock in {}ms, cost: {}ms, operation: scheduled_upload", TimeUnit.NANOSECONDS.toMillis(DEFAULT_LOCK_WARNING_TIMEOUT), TimeUnit.NANOSECONDS.toMillis(time.nanoseconds() - startTime)); - } - - if (System.currentTimeMillis() - lastUploadTimestamp >= config.batchInterval()) { - unsafeUpload(false); - } - } catch (Throwable ignore) { - } finally { - lock.writeLock().unlock(); - } - }, config.batchInterval(), config.batchInterval(), TimeUnit.MILLISECONDS); - - monitorService.scheduleWithFixedDelay(() -> { - try { - long count = pendingFutureMap.values() - .stream() - .filter(uploadTime -> time.nanoseconds() - uploadTime > DEFAULT_UPLOAD_WARNING_TIMEOUT) - .count(); - if (count > 0) { - log.error("Found {} pending upload tasks exceed 5s.", count); - } - } catch (Throwable ignore) { - } - }, 1, 1, TimeUnit.SECONDS); - - closed = false; - } - - @Override - public void close() { - closed = true; - - if (executorService != null && !executorService.isShutdown()) { - executorService.shutdown(); - try { - if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) { - log.error("Main executor {} did not terminate in time", executorService); - executorService.shutdownNow(); - } - } catch (InterruptedException e) { - log.error("Failed to shutdown main executor service.", e); - executorService.shutdownNow(); - } - } - - lock.writeLock().lock(); - try { - unsafeUpload(true); - } catch (Throwable throwable) { - log.error("Failed to flush records when close.", throwable); - } finally { - lock.writeLock().unlock(); - } - - // Wait for all upload tasks to complete. - if (!pendingFutureMap.isEmpty()) { - log.info("Wait for {} pending upload tasks to complete.", pendingFutureMap.size()); - CompletableFuture.allOf(pendingFutureMap.keySet().toArray(new CompletableFuture[0])).join(); - } - - if (monitorService != null && !monitorService.isShutdown()) { - monitorService.shutdown(); - try { - if (!monitorService.awaitTermination(30, TimeUnit.SECONDS)) { - log.error("Monitor executor {} did not terminate in time", executorService); - monitorService.shutdownNow(); - } - } catch (InterruptedException e) { - log.error("Failed to shutdown monitor executor service.", e); - monitorService.shutdownNow(); - } - } - - if (callbackService != null && !callbackService.isShutdown()) { - callbackService.shutdown(); - try { - if (!callbackService.awaitTermination(30, TimeUnit.SECONDS)) { - log.error("Callback executor {} did not terminate in time", executorService); - callbackService.shutdownNow(); - } - } catch (InterruptedException e) { - log.error("Failed to shutdown callback executor service.", e); - callbackService.shutdownNow(); - } - } - - log.info("S3 WAL record accumulator is closed."); - } - - public long nextOffset() { - return nextOffset.get(); - } - - public long flushedOffset() { - return flushedOffset.get(); - } - - public List objectList() throws WALFencedException { - checkStatus(); - - List list = new ArrayList<>(objectMap.size() + previousObjectMap.size()); - list.addAll(previousObjectMap.values()); - list.addAll(objectMap.values()); - return list; - } - - // Visible for testing - public ScheduledExecutorService executorService() { - return executorService; - } - - protected void checkStatus() throws WALFencedException { - if (closed) { - throw new IllegalStateException("WAL is closed."); - } - - if (fenced) { - throw new WALFencedException("WAL is fenced."); - } - } - - protected void checkWriteStatus() throws WALFencedException { - if (config.failover()) { - throw new IllegalStateException("WAL is in failover mode."); - } - - checkStatus(); - } - - private boolean shouldUpload() { - Record firstRecord = bufferQueue.peekFirst(); - if (firstRecord == null || uploadMap.size() >= config.maxInflightUploadCount()) { - return false; - } - - return System.currentTimeMillis() - lastUploadTimestamp >= config.batchInterval() - || nextOffset.get() - firstRecord.offset > config.maxBytesInBatch(); - } - - public long append(long recordSize, Function recordSupplier, - CompletableFuture future) throws OverCapacityException, WALFencedException { - long startTime = time.nanoseconds(); - checkWriteStatus(); - - // Check if there is too much data in the S3 WAL. - if (nextOffset.get() - flushedOffset.get() > config.maxUnflushedBytes()) { - throw new OverCapacityException("Too many unflushed bytes.", true); - } - - if (objectMap.size() + config.maxInflightUploadCount() >= 3000) { - throw new OverCapacityException("Too many WAL objects.", false); - } - - if (shouldUpload() && lock.writeLock().tryLock()) { - try { - if (time.nanoseconds() - startTime > DEFAULT_LOCK_WARNING_TIMEOUT) { - log.warn("Failed to acquire lock in {}ms, cost: {}ms, operation: append_upload", TimeUnit.NANOSECONDS.toMillis(DEFAULT_LOCK_WARNING_TIMEOUT), TimeUnit.NANOSECONDS.toMillis(time.nanoseconds() - startTime)); - } - - if (shouldUpload()) { - unsafeUpload(false); - } - } finally { - lock.writeLock().unlock(); - } - } - - long acquireAppendLockTime = time.nanoseconds(); - lock.readLock().lock(); - try { - if (time.nanoseconds() - acquireAppendLockTime > DEFAULT_LOCK_WARNING_TIMEOUT) { - log.warn("Failed to acquire lock in {}ms, cost: {}ms, operation: append", TimeUnit.NANOSECONDS.toMillis(DEFAULT_LOCK_WARNING_TIMEOUT), TimeUnit.NANOSECONDS.toMillis(time.nanoseconds() - acquireAppendLockTime)); - } - - long offset = nextOffset.getAndAdd(recordSize); - future.whenComplete((v, throwable) -> { - if (throwable != null) { - log.error("Failed to append record to S3 WAL: {}", offset, throwable); - } else { - ObjectWALMetricsManager.recordOperationDataSize(recordSize, "append"); - } - ObjectWALMetricsManager.recordOperationLatency(time.nanoseconds() - acquireAppendLockTime, "append", throwable == null); - }); - - bufferQueue.offer(new Record(offset, recordSupplier.apply(offset), future)); - bufferedDataBytes.addAndGet(recordSize); - - return offset; - } finally { - lock.readLock().unlock(); - } - } - - public CompletableFuture reset() throws WALFencedException { - checkStatus(); - - if (objectMap.isEmpty() && previousObjectMap.isEmpty()) { - return CompletableFuture.completedFuture(null); - } - - long startTime = time.nanoseconds(); - AtomicLong deletedObjectSize = new AtomicLong(); - List deleteObjectList = new ArrayList<>(); - - previousObjectMap.forEach((k, v) -> { - deleteObjectList.add(new ObjectStorage.ObjectPath(v.bucketId(), v.path())); - deletedObjectSize.addAndGet(v.length()); - previousObjectMap.remove(k); - }); - objectMap.forEach((k, v) -> { - deleteObjectList.add(new ObjectStorage.ObjectPath(v.bucketId(), v.path())); - deletedObjectSize.addAndGet(v.length()); - objectMap.remove(k); - }); - - return objectStorage.delete(deleteObjectList) - .whenComplete((v, throwable) -> { - ObjectWALMetricsManager.recordOperationLatency(time.nanoseconds() - startTime, "reset", throwable == null); - objectDataBytes.addAndGet(-1 * deletedObjectSize.get()); - - // Never fail the delete task, the under layer storage will retry forever. - if (throwable != null) { - log.error("Failed to delete objects when trim S3 WAL: {}", deleteObjectList, throwable); - } - }); - } - - // Trim objects where the last offset is less than or equal to the given offset. - public CompletableFuture trim(long offset) throws WALFencedException { - checkStatus(); - - if (objectMap.isEmpty() || offset < objectMap.firstKey() || offset > flushedOffset.get()) { - return CompletableFuture.completedFuture(null); - } - - long startTime = time.nanoseconds(); - - List deleteObjectList = new ArrayList<>(); - AtomicLong deletedObjectSize = new AtomicLong(); - - objectMap.headMap(offset, true) - .forEach((k, v) -> { - deleteObjectList.add(new ObjectStorage.ObjectPath(v.bucketId(), v.path())); - deletedObjectSize.addAndGet(v.length()); - objectMap.remove(k); - }); - - if (deleteObjectList.isEmpty()) { - return CompletableFuture.completedFuture(null); - } - - return objectStorage.delete(deleteObjectList) - .whenComplete((v, throwable) -> { - ObjectWALMetricsManager.recordOperationLatency(time.nanoseconds() - startTime, "trim", throwable == null); - objectDataBytes.addAndGet(-1 * deletedObjectSize.get()); - - // Never fail the delete task, the under layer storage will retry forever. - if (throwable != null) { - log.error("Failed to delete objects when trim S3 WAL: {}", deleteObjectList, throwable); - } - }); - } - - // Not thread safe, caller should hold lock. - // Visible for testing. - public void unsafeUpload(boolean force) throws WALFencedException { - if (!force) { - checkWriteStatus(); - } - - if (bufferQueue.isEmpty()) { - return; - } - - int size = bufferQueue.size(); - PriorityQueue recordQueue = new PriorityQueue<>(size, Comparator.comparingLong(o -> o.offset)); - - for (int i = 0; i < size; i++) { - Record record = bufferQueue.poll(); - if (record != null) { - recordQueue.offer(record); - } - } - - // Trigger upload until the buffer is empty. - while (!recordQueue.isEmpty()) { - unsafeUpload(recordQueue); - } - } - - // Not thread safe, caller should hold lock. - private void unsafeUpload(PriorityQueue recordQueue) { - long startTime = time.nanoseconds(); - - // Build data buffer. - CompositeByteBuf dataBuffer = ByteBufAlloc.compositeByteBuffer(); - List recordList = new LinkedList<>(); - - long stickyRecordLength = 0; - if (!recordQueue.isEmpty()) { - Record firstRecord = recordQueue.peek(); - if (firstRecord.record.readerIndex() != 0) { - stickyRecordLength = firstRecord.record.readableBytes(); - } - } - - while (!recordQueue.isEmpty()) { - // The retained bytes in the batch must larger than record header size. - long retainedBytesInBatch = config.maxBytesInBatch() - dataBuffer.readableBytes() - WALObjectHeader.WAL_HEADER_SIZE; - if (config.strictBatchLimit() && retainedBytesInBatch <= RecordHeader.RECORD_HEADER_SIZE) { - break; - } - - Record record = recordQueue.poll(); - - // Records larger than the batch size will be uploaded immediately. - assert record != null; - if (config.strictBatchLimit() && record.record.readableBytes() >= config.maxBytesInBatch() - WALObjectHeader.WAL_HEADER_SIZE) { - dataBuffer.addComponent(true, record.record); - recordList.add(record); - break; - } - - if (config.strictBatchLimit() && record.record.readableBytes() > retainedBytesInBatch) { - // The records will be split into multiple objects. - ByteBuf slice = record.record.retainedSlice(0, (int) retainedBytesInBatch).asReadOnly(); - dataBuffer.addComponent(true, slice); - - // Update the record buffer and offset. - record.record.skipBytes((int) retainedBytesInBatch); - record.offset += retainedBytesInBatch; - recordQueue.offer(record); - break; - } - - dataBuffer.addComponent(true, record.record); - recordList.add(record); - } - - // Build object buffer. - long firstOffset = recordList.get(0).offset; - long dataLength = dataBuffer.readableBytes(); - // Exclusive end offset - long endOffset = firstOffset + dataLength; - - CompositeByteBuf objectBuffer = ByteBufAlloc.compositeByteBuffer(); - WALObjectHeader header = new WALObjectHeader(firstOffset, dataLength, stickyRecordLength, config.nodeId(), config.epoch()); - objectBuffer.addComponent(true, header.marshal()); - objectBuffer.addComponent(true, dataBuffer); - - // Trigger upload. - int objectLength = objectBuffer.readableBytes(); - uploadMap.put(firstOffset, recordList); - - // Enable fast retry. - ObjectStorage.WriteOptions writeOptions = new ObjectStorage.WriteOptions().enableFastRetry(true); - String path = objectPrefix + firstOffset; - CompletableFuture uploadFuture = objectStorage.write(writeOptions, path, objectBuffer); - - CompletableFuture finalFuture = recordUploadMetrics(uploadFuture, startTime, objectLength) - .thenAccept(result -> { - long lockStartTime = time.nanoseconds(); - lock.writeLock().lock(); - try { - if (time.nanoseconds() - lockStartTime > DEFAULT_LOCK_WARNING_TIMEOUT) { - log.warn("Failed to acquire lock in {}ms, cost: {}ms, operation: upload", TimeUnit.NANOSECONDS.toMillis(DEFAULT_LOCK_WARNING_TIMEOUT), TimeUnit.NANOSECONDS.toMillis(time.nanoseconds() - lockStartTime)); - } - - objectMap.put(endOffset, new WALObject(result.bucket(), path, firstOffset, objectLength)); - objectDataBytes.addAndGet(objectLength); - - List uploadedRecords = uploadMap.remove(firstOffset); - - // Update flushed offset - if (!uploadMap.isEmpty()) { - flushedOffset.set(uploadMap.firstKey()); - } else if (!bufferQueue.isEmpty()) { - flushedOffset.set(bufferQueue.getFirst().offset); - } else { - flushedOffset.set(nextOffset.get()); - } - - // Release lock and complete future in callback thread. - callbackService.submit(() -> uploadedRecords.forEach(record -> record.future.complete(flushedOffset::get))); - } finally { - lock.writeLock().unlock(); - } - }) - .whenComplete((v, throwable) -> { - bufferedDataBytes.addAndGet(-dataLength); - throwable = ExceptionUtils.getRootCause(throwable); - if (throwable instanceof WALFencedException) { - List uploadedRecords = uploadMap.remove(firstOffset); - Throwable finalThrowable = throwable; - // Release lock and complete future in callback thread. - callbackService.submit(() -> uploadedRecords.forEach(record -> record.future.completeExceptionally(finalThrowable))); - } else if (throwable != null) { - // Never fail the write task, the under layer storage will retry forever. - log.error("[Bug] Failed to write records to S3: {}", firstOffset, throwable); - } - }); - - pendingFutureMap.put(finalFuture, time.nanoseconds()); - finalFuture.whenComplete((v, throwable) -> pendingFutureMap.remove(finalFuture)); - lastUploadTimestamp = System.currentTimeMillis(); - } - - protected CompletableFuture recordUploadMetrics( - CompletableFuture future, - long startTime, long objectLength) { - return future.whenComplete((result, throwable) -> { - ObjectWALMetricsManager.recordOperationLatency(time.nanoseconds() - startTime, "upload", throwable == null); - ObjectWALMetricsManager.recordOperationDataSize(objectLength, "upload"); - }); - } - - protected static class Record { - public final ByteBuf record; - public final CompletableFuture future; - public long offset; - - public Record(long offset, ByteBuf record, - CompletableFuture future) { - this.offset = offset; - this.record = record; - this.future = future; - } - } - - public static class WALObject implements Comparable { - private final short bucketId; - private final String path; - private final long startOffset; - private final long length; - - public WALObject(short bucketId, String path, long startOffset, long length) { - this.bucketId = bucketId; - this.path = path; - this.startOffset = startOffset; - this.length = length; - } - - @Override - public int compareTo(WALObject o) { - return Long.compare(startOffset, o.startOffset); - } - - public short bucketId() { - return bucketId; - } - - public String path() { - return path; - } - - public long startOffset() { - return startOffset; - } - - public long length() { - return length; - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecoverIterator.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecoverIterator.java new file mode 100644 index 0000000000..58799fa7e4 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/RecoverIterator.java @@ -0,0 +1,277 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.network.ThrottleStrategy; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.wal.RecoverResult; +import com.automq.stream.s3.wal.common.RecordHeader; +import com.automq.stream.s3.wal.common.RecoverResultImpl; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; +import com.automq.stream.s3.wal.util.WALUtil; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Queue; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; +import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_WITHOUT_CRC_SIZE; + +public class RecoverIterator implements Iterator { + private static final Logger LOGGER = LoggerFactory.getLogger(RecoverIterator.class); + private final ObjectStorage objectStorage; + private final int readAheadObjectSize; + + private final long trimOffset; + private final List objectList; + private final Queue> readAheadQueue; + private final TreeMap startOffset2Epoch = new TreeMap<>(); + + private RecoverResult nextRecord = null; + private int nextIndex = 0; + private ByteBuf dataBuffer = Unpooled.EMPTY_BUFFER; + + public RecoverIterator(List objectList, ObjectStorage objectStorage, + int readAheadObjectSize) { + this.trimOffset = getTrimOffset(objectList, objectStorage); + this.objectList = getContinuousFromTrimOffset(objectList, trimOffset); + this.objectStorage = objectStorage; + this.readAheadObjectSize = readAheadObjectSize; + this.readAheadQueue = new ArrayDeque<>(readAheadObjectSize); + + long lastEpoch = -1L; + for (WALObject object : objectList) { + if (object.epoch() != lastEpoch) { + startOffset2Epoch.put(object.startOffset(), object.epoch()); + lastEpoch = object.epoch(); + } + } + + // Fill the read ahead queue. + for (int i = 0; i < readAheadObjectSize; i++) { + tryReadAhead(); + } + } + + /** + * Get the latest trim offset from the newest object. + */ + private static long getTrimOffset(List objectList, ObjectStorage objectStorage) { + if (objectList.isEmpty()) { + return -1; + } + + WALObject object = objectList.get(objectList.size() - 1); + ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions() + .throttleStrategy(ThrottleStrategy.BYPASS) + .bucket(object.bucketId()); + ByteBuf buffer = objectStorage.rangeRead(options, object.path(), 0, Math.min(WALObjectHeader.MAX_WAL_HEADER_SIZE, object.length())).join(); + WALObjectHeader header = WALObjectHeader.unmarshal(buffer); + buffer.release(); + return header.trimOffset(); + } + + // Visible for testing. + static List getContinuousFromTrimOffset(List objectList, long trimOffset) { + if (objectList.isEmpty()) { + return Collections.emptyList(); + } + + int startIndex = objectList.size(); + for (int i = 0; i < objectList.size(); i++) { + if (objectList.get(i).endOffset() > trimOffset) { + startIndex = i; + break; + } + } + if (startIndex > 0) { + for (int i = 0; i < startIndex; i++) { + LOGGER.info("drop trimmed object: {}", objectList.get(i)); + } + } + if (startIndex >= objectList.size()) { + return Collections.emptyList(); + } + + int endIndex = startIndex + 1; + for (int i = startIndex + 1; i < objectList.size(); i++) { + if (objectList.get(i).startOffset() != objectList.get(i - 1).endOffset()) { + break; + } + endIndex = i + 1; + } + if (endIndex < objectList.size()) { + for (int i = endIndex; i < objectList.size(); i++) { + LOGGER.warn("drop discontinuous object: {}", objectList.get(i)); + } + } + + return new ArrayList<>(objectList.subList(startIndex, endIndex)); + } + + @Override + public boolean hasNext() { + if (nextRecord != null) { + return true; + } else { + while (hasNext0()) { + RecoverResult record = next0(); + DefaultRecordOffset recordOffset = (DefaultRecordOffset) record.recordOffset(); + //noinspection DataFlowIssue + if ((recordOffset.offset() <= trimOffset) + || (record.record().getStreamId() == -1L && record.record().getEpoch() == -1L)) { + record.record().release(); + continue; + } + nextRecord = record; + return true; + } + return false; + } + } + + private boolean hasNext0() { + return dataBuffer.isReadable() || !readAheadQueue.isEmpty() || nextIndex < objectList.size(); + } + + private void loadNextBuffer() { + // Please call hasNext() before calling loadNextBuffer(). + byte[] buffer = Objects.requireNonNull(readAheadQueue.poll()).join(); + dataBuffer = Unpooled.wrappedBuffer(buffer); + + // Check header + WALObjectHeader header = WALObjectHeader.unmarshal(dataBuffer); + dataBuffer.skipBytes(header.size()); + } + + private void tryReadAhead() { + if (readAheadQueue.size() < readAheadObjectSize && nextIndex < objectList.size()) { + WALObject object = objectList.get(nextIndex++); + ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions().throttleStrategy(ThrottleStrategy.BYPASS).bucket(object.bucketId()); + CompletableFuture readFuture = objectStorage.rangeRead(options, object.path(), 0, object.length()) + .thenApply(buffer -> { + // Copy the result buffer and release it. + byte[] bytes = new byte[buffer.readableBytes()]; + buffer.readBytes(bytes); + buffer.release(); + return bytes; + }); + readAheadQueue.add(readFuture); + } + } + + @Override + public RecoverResult next() { + if (nextRecord != null || hasNext()) { + // - If the nextRecord is already read ahead. + // - Or #hasNext() is true, it means the nextRecord is already ready. + RecoverResult rst = nextRecord; + nextRecord = null; + return rst; + } else { + return null; + } + } + + public RecoverResult next0() { + // If there is no more data to read, return null. + if (!dataBuffer.isReadable()) { + loadNextBuffer(); + } + + // TODO: simple the code without strict batch + + // Try to read next object. + tryReadAhead(); + + ByteBuf recordHeaderBuf = dataBuffer.readBytes(RECORD_HEADER_SIZE); + RecordHeader header = new RecordHeader(recordHeaderBuf); + + if (header.getRecordHeaderCRC() != WALUtil.crc32(recordHeaderBuf, RECORD_HEADER_WITHOUT_CRC_SIZE)) { + recordHeaderBuf.release(); + throw new IllegalStateException("Record header crc check failed."); + } + recordHeaderBuf.release(); + + if (header.getMagicCode() != RecordHeader.RECORD_HEADER_DATA_MAGIC_CODE) { + throw new IllegalStateException("Invalid magic code in record header."); + } + + int length = header.getRecordBodyLength(); + + ByteBuf recordBuf = ByteBufAlloc.byteBuffer(length); + if (dataBuffer.readableBytes() < length) { + // Read the remain data and release the buffer. + dataBuffer.readBytes(recordBuf, dataBuffer.readableBytes()); + dataBuffer.release(); + + // Read from next buffer. + if (!hasNext()) { + throw new IllegalStateException("[Bug] There is a record part but no more data to read."); + } + loadNextBuffer(); + dataBuffer.readBytes(recordBuf, length - recordBuf.readableBytes()); + } else { + dataBuffer.readBytes(recordBuf, length); + } + + if (!dataBuffer.isReadable()) { + dataBuffer.release(); + } + + if (header.getRecordBodyCRC() != WALUtil.crc32(recordBuf)) { + recordBuf.release(); + throw new IllegalStateException("Record body crc check failed."); + } + + long offset = header.getRecordBodyOffset() - RECORD_HEADER_SIZE; + int size = recordBuf.readableBytes() + RECORD_HEADER_SIZE; + + try { + return new RecoverResultImpl(StreamRecordBatch.parse(recordBuf, false), DefaultRecordOffset.of(getEpoch(offset), offset, size)); + } finally { + recordBuf.release(); + } + } + + private long getEpoch(long offset) { + Map.Entry entry = startOffset2Epoch.floorEntry(offset); + if (entry == null) { + LOGGER.error("[BUG] Cannot find any epoch for offset {}, startOffset2epoch={}", offset, startOffset2Epoch); + throw new IllegalStateException("[BUG] Cannot find any epoch for offset"); + } + return entry.getValue(); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObject.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObject.java new file mode 100644 index 0000000000..eb947eebaa --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObject.java @@ -0,0 +1,102 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import java.util.Objects; + +public class WALObject implements Comparable { + private final short bucketId; + private final String path; + private final long epoch; + private final long startOffset; + private final long endOffset; + private final long length; + + public WALObject(short bucketId, String path, long epoch, long startOffset, long length) { + this.bucketId = bucketId; + this.path = path; + this.epoch = epoch; + this.startOffset = startOffset; + this.endOffset = WALObjectHeader.calculateEndOffsetV0(startOffset, length); + this.length = length; + } + + public WALObject(short bucketId, String path, long epoch, long startOffset, long endOffset, long length) { + this.bucketId = bucketId; + this.path = path; + this.epoch = epoch; + this.startOffset = startOffset; + this.endOffset = endOffset; + this.length = length; + } + + @Override + public int compareTo(WALObject o) { + return Long.compare(startOffset, o.startOffset); + } + + public short bucketId() { + return bucketId; + } + + public String path() { + return path; + } + + public long epoch() { + return epoch; + } + + public long startOffset() { + return startOffset; + } + + public long length() { + return length; + } + + public long endOffset() { + return endOffset; + } + + @Override + public String toString() { + return "WALObject{" + + "bucketId=" + bucketId + + ", path='" + path + '\'' + + ", startOffset=" + startOffset + + ", endOffset=" + endOffset + + ", length=" + length + + '}'; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof WALObject)) + return false; + WALObject object = (WALObject) o; + return bucketId == object.bucketId && startOffset == object.startOffset && endOffset == object.endOffset && length == object.length && Objects.equals(path, object.path); + } + + @Override + public int hashCode() { + return Objects.hash(bucketId, path, startOffset, endOffset, length); + } +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObjectHeader.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObjectHeader.java index 0b54295903..ab73e40e46 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObjectHeader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/WALObjectHeader.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.wal.impl.object; @@ -14,68 +22,138 @@ import com.automq.stream.s3.ByteBufAlloc; import com.automq.stream.s3.wal.exception.UnmarshalException; +import java.util.Map; +import java.util.Objects; + import io.netty.buffer.ByteBuf; public class WALObjectHeader { - public static final int WAL_HEADER_MAGIC_CODE = 0x12345678; - public static final int WAL_HEADER_SIZE = 4 // magic code - + 8 // start offset - + 8 // body length - + 8 // sticky record length - + 4 // node id - + 8; // node epoch - - private int magicCode0 = WAL_HEADER_MAGIC_CODE; - private long startOffset1; - private long length2; - private long stickyRecordLength3; - private int nodeId4; - private long epoch5; + // Visible for testing. + static final int WAL_HEADER_MAGIC_CODE_V0 = 0x12345678; + static final int WAL_HEADER_SIZE_V0 = 4 // magic code + + 8 // start offset + + 8 // body length + + 8 // sticky record length + + 4 // node id + + 8; // node epoch + static final int WAL_HEADER_MAGIC_CODE_V1 = 0xEDCBA987; + static final int WAL_HEADER_SIZE_V1 = WAL_HEADER_SIZE_V0 + + 8; // trim offset + static final int MAX_WAL_HEADER_SIZE = Math.max(WAL_HEADER_SIZE_V0, WAL_HEADER_SIZE_V1); + + private static final Map WAL_HEADER_SIZES = Map.of( + WAL_HEADER_MAGIC_CODE_V0, WAL_HEADER_SIZE_V0, + WAL_HEADER_MAGIC_CODE_V1, WAL_HEADER_SIZE_V1 + ); + + public static final int DEFAULT_WAL_MAGIC_CODE = WAL_HEADER_MAGIC_CODE_V1; + public static final int DEFAULT_WAL_HEADER_SIZE = WAL_HEADER_SIZE_V1; + + private final int magicCode0; + private final long startOffset1; + private final long length2; + @Deprecated + private final long stickyRecordLength3; + private final int nodeId4; + private final long epoch5; + private final long trimOffset6; - public WALObjectHeader() { + public WALObjectHeader(long startOffset, long length, long stickyRecordLength, int nodeId, long epoch) { + this.magicCode0 = WAL_HEADER_MAGIC_CODE_V0; + this.startOffset1 = startOffset; + this.length2 = length; + this.stickyRecordLength3 = stickyRecordLength; + this.nodeId4 = nodeId; + this.epoch5 = epoch; + this.trimOffset6 = -1; } - public WALObjectHeader(long startOffset, long length, long stickyRecordLength, int nodeId, long epoch) { + public WALObjectHeader(long startOffset, long length, long stickyRecordLength, int nodeId, long epoch, long trimOffset) { + this.magicCode0 = WAL_HEADER_MAGIC_CODE_V1; this.startOffset1 = startOffset; this.length2 = length; this.stickyRecordLength3 = stickyRecordLength; this.nodeId4 = nodeId; this.epoch5 = epoch; + this.trimOffset6 = trimOffset; + } + + /** + * In the historical version V0, the endOffset of each WAL object is calculated directly from the path and size of the object. + * This method is used to be compatible with this case. + * + * @param startOffset the start offset of the WAL object, get from the path + * @param length the size of the WAL object + * @return the end offset of the WAL object + */ + public static long calculateEndOffsetV0(long startOffset, long length) { + return startOffset + length - WAL_HEADER_SIZE_V0; } public static WALObjectHeader unmarshal(ByteBuf buf) throws UnmarshalException { - if (buf.readableBytes() < WAL_HEADER_SIZE) { - throw new UnmarshalException(String.format("WALHeader does not have enough bytes, Recovered: [%d] expect: [%d]", buf.readableBytes(), WAL_HEADER_SIZE)); + buf.markReaderIndex(); + + int size = buf.readableBytes(); + if (size < 4) { + throw new UnmarshalException(String.format("Insufficient bytes to read magic code, Recovered: [%d] expect: [%d]", size, 4)); } - WALObjectHeader walObjectHeader = new WALObjectHeader(); - buf.markReaderIndex(); - walObjectHeader.magicCode0 = buf.readInt(); - if (walObjectHeader.magicCode0 != WAL_HEADER_MAGIC_CODE) { - throw new UnmarshalException(String.format("WALHeader magic code not match, Recovered: [%d] expect: [%d]", walObjectHeader.magicCode0, WAL_HEADER_MAGIC_CODE)); + int magicCode = buf.readInt(); + if (!WAL_HEADER_SIZES.containsKey(magicCode)) { + throw new UnmarshalException(String.format("WALHeader magic code not match, Recovered: [%d] expect: [%s]", magicCode, WAL_HEADER_SIZES.keySet())); + } + if (size < WAL_HEADER_SIZES.get(magicCode)) { + throw new UnmarshalException(String.format("WALHeader does not have enough bytes, Recovered: [%d] expect: [%d]", size, WAL_HEADER_SIZES.get(magicCode))); } - walObjectHeader.startOffset1 = buf.readLong(); - walObjectHeader.length2 = buf.readLong(); - walObjectHeader.stickyRecordLength3 = buf.readLong(); - walObjectHeader.nodeId4 = buf.readInt(); - walObjectHeader.epoch5 = buf.readLong(); - buf.resetReaderIndex(); + WALObjectHeader header = null; + if (magicCode == WAL_HEADER_MAGIC_CODE_V1) { + header = new WALObjectHeader(buf.readLong(), buf.readLong(), buf.readLong(), buf.readInt(), buf.readLong(), buf.readLong()); + } else if (magicCode == WAL_HEADER_MAGIC_CODE_V0) { + header = new WALObjectHeader(buf.readLong(), buf.readLong(), buf.readLong(), buf.readInt(), buf.readLong()); + } - return walObjectHeader; + buf.resetReaderIndex(); + return header; } public ByteBuf marshal() { - ByteBuf buf = ByteBufAlloc.byteBuffer(WAL_HEADER_SIZE); + if (magicCode0 == WAL_HEADER_MAGIC_CODE_V1) { + return marshalV1(); + } else if (magicCode0 == WAL_HEADER_MAGIC_CODE_V0) { + return marshalV0(); + } else { + throw new IllegalStateException("Invalid magic code: " + magicCode0); + } + } + + private ByteBuf marshalV0() { + ByteBuf buf = ByteBufAlloc.byteBuffer(WAL_HEADER_SIZE_V0); + buf.writeInt(magicCode0); + buf.writeLong(startOffset1); + buf.writeLong(length2); + buf.writeLong(stickyRecordLength3); + buf.writeInt(nodeId4); + buf.writeLong(epoch5); + return buf; + } + + private ByteBuf marshalV1() { + ByteBuf buf = ByteBufAlloc.byteBuffer(WAL_HEADER_SIZE_V1); buf.writeInt(magicCode0); buf.writeLong(startOffset1); buf.writeLong(length2); buf.writeLong(stickyRecordLength3); buf.writeInt(nodeId4); buf.writeLong(epoch5); + buf.writeLong(trimOffset6); return buf; } + public int size() { + return WAL_HEADER_SIZES.get(magicCode0); + } + public int magicCode() { return magicCode0; } @@ -88,10 +166,6 @@ public long length() { return length2; } - public long stickyRecordLength() { - return stickyRecordLength3; - } - public int nodeId() { return nodeId4; } @@ -99,4 +173,21 @@ public int nodeId() { public long epoch() { return epoch5; } + + public long trimOffset() { + return trimOffset6; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof WALObjectHeader)) + return false; + WALObjectHeader header = (WALObjectHeader) o; + return magicCode0 == header.magicCode0 && startOffset1 == header.startOffset1 && length2 == header.length2 && stickyRecordLength3 == header.stickyRecordLength3 && nodeId4 == header.nodeId4 && epoch5 == header.epoch5 && trimOffset6 == header.trimOffset6; + } + + @Override + public int hashCode() { + return Objects.hash(magicCode0, startOffset1, length2, stickyRecordLength3, nodeId4, epoch5, trimOffset6); + } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/Writer.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/Writer.java new file mode 100644 index 0000000000..abc9279594 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/Writer.java @@ -0,0 +1,46 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.RecordOffset; +import com.automq.stream.s3.wal.RecoverResult; +import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.s3.wal.exception.WALFencedException; + +import java.util.Iterator; +import java.util.concurrent.CompletableFuture; + +public interface Writer { + void start(); + + void close(); + + CompletableFuture append(StreamRecordBatch streamRecordBatch) throws OverCapacityException; + + RecordOffset confirmOffset(); + + CompletableFuture reset() throws WALFencedException; + + CompletableFuture trim(RecordOffset recordOffset) throws WALFencedException; + + Iterator recover(); +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/metrics/ObjectWALMetricsManager.java b/s3stream/src/main/java/com/automq/stream/s3/wal/metrics/ObjectWALMetricsManager.java deleted file mode 100644 index 981bb22093..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/metrics/ObjectWALMetricsManager.java +++ /dev/null @@ -1,95 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.metrics; - -import com.automq.stream.s3.metrics.NoopLongHistogram; -import com.automq.stream.s3.metrics.NoopObservableLongGauge; - -import java.util.List; -import java.util.concurrent.TimeUnit; -import java.util.function.Supplier; - -import io.opentelemetry.api.common.Attributes; -import io.opentelemetry.api.metrics.LongHistogram; -import io.opentelemetry.api.metrics.Meter; -import io.opentelemetry.api.metrics.ObservableLongGauge; - -public class ObjectWALMetricsManager { - private static ObservableLongGauge inflightuploadCount = new NoopObservableLongGauge(); - private static ObservableLongGauge bufferedDataSizeInBytes = new NoopObservableLongGauge(); - private static ObservableLongGauge objectDataSizeInBytes = new NoopObservableLongGauge(); - private static LongHistogram operationLatencyInMillis = new NoopLongHistogram(); - private static LongHistogram operationDataSizeInBytes = new NoopLongHistogram(); - - private static Supplier inflightUploadCountSupplier = () -> 0L; - private static Supplier bufferedDataInBytesSupplier = () -> 0L; - private static Supplier objectDataInBytesSupplier = () -> 0L; - - public static void initMetrics(Meter meter) { - initMetrics(meter, ""); - } - - public static void initMetrics(Meter meter, String prefix) { - operationLatencyInMillis = meter - .histogramBuilder(prefix + "operation_latency") - .ofLongs() - .setUnit("milliseconds") - .setExplicitBucketBoundariesAdvice(List.of(10L, 100L, 150L, 200L, 300L, 400L, 500L, 750L, 1000L, 3 * 1000L)) - .setDescription("Operation latency in milliseconds") - .build(); - operationDataSizeInBytes = meter - .histogramBuilder(prefix + "data_size") - .ofLongs() - .setUnit("bytes") - .setDescription("Operation size in bytes") - .setExplicitBucketBoundariesAdvice(List.of(512L, 1024L, 16 * 1024L, 32 * 1024L, 64 * 1024L, 128 * 1024L, 256 * 1024L, 512 * 1024L, - 1024 * 1024L, 4 * 1024L * 1024L, 8 * 1024L * 1024L, 16 * 1024L * 1024L)) - .build(); - inflightuploadCount = meter - .gaugeBuilder(prefix + "inflight_upload_count") - .setDescription("Inflight upload count") - .ofLongs() - .buildWithCallback(measurement -> measurement.record(inflightUploadCountSupplier.get(), Attributes.empty())); - bufferedDataSizeInBytes = meter - .gaugeBuilder(prefix + "buffered_data_size") - .setDescription("Buffered data size") - .ofLongs() - .buildWithCallback(measurement -> measurement.record(bufferedDataInBytesSupplier.get(), Attributes.empty())); - objectDataSizeInBytes = meter - .gaugeBuilder(prefix + "object_data_size") - .setDescription("Object data size in S3") - .ofLongs() - .buildWithCallback(measurement -> measurement.record(objectDataInBytesSupplier.get(), Attributes.empty())); - } - - public static void recordOperationLatency(long latencyInNanos, String operation, boolean success) { - Attributes attributes = Attributes.builder().put("operation", operation).put("success", success).build(); - operationLatencyInMillis.record(TimeUnit.NANOSECONDS.toMillis(latencyInNanos), attributes); - } - - public static void recordOperationDataSize(long size, String operation) { - Attributes attributes = Attributes.builder().put("operation", operation).build(); - operationDataSizeInBytes.record(size, attributes); - } - - public static void setInflightUploadCountSupplier(Supplier inflightuploadCountSupplier) { - ObjectWALMetricsManager.inflightUploadCountSupplier = inflightuploadCountSupplier; - } - - public static void setBufferedDataInBytesSupplier(Supplier bufferedDataInBytesSupplier) { - ObjectWALMetricsManager.bufferedDataInBytesSupplier = bufferedDataInBytesSupplier; - } - - public static void setObjectDataInBytesSupplier(Supplier objectDataInBytesSupplier) { - ObjectWALMetricsManager.objectDataInBytesSupplier = objectDataInBytesSupplier; - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/AbstractWALChannel.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/AbstractWALChannel.java deleted file mode 100644 index 870c2404cc..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/AbstractWALChannel.java +++ /dev/null @@ -1,124 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.utils.Threads; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.IOException; -import java.util.concurrent.TimeUnit; - -import io.netty.buffer.ByteBuf; - -public abstract class AbstractWALChannel implements WALChannel { - - private static final Logger LOGGER = LoggerFactory.getLogger(AbstractWALChannel.class); - - /** - * Flag to indicate if the WAL has failed. - * It will be set to true if an IO operation fails continuously, and it will never be reset. - * Any IO operation will fail immediately if this flag is true. - */ - private volatile boolean failed = false; - - @Override - public void write(ByteBuf src, long position) throws IOException { - checkFailed(); - doWrite(src, position); - } - - @Override - public void retryWrite(ByteBuf src, long position, long retryIntervalMillis, - long retryTimeoutMillis) throws IOException { - checkFailed(); - retry(() -> write(src, position), retryIntervalMillis, retryTimeoutMillis); - } - - @Override - public void flush() throws IOException { - checkFailed(); - doFlush(); - } - - @Override - public void retryFlush(long retryIntervalMillis, long retryTimeoutMillis) throws IOException { - checkFailed(); - retry(this::flush, retryIntervalMillis, retryTimeoutMillis); - } - - @Override - public int read(ByteBuf dst, long position, int length) throws IOException { - checkFailed(); - return doRead(dst, position, length); - } - - @Override - public int retryRead(ByteBuf dst, long position, int length, long retryIntervalMillis, - long retryTimeoutMillis) throws IOException { - checkFailed(); - return retry(() -> read(dst, position, length), retryIntervalMillis, retryTimeoutMillis); - } - - private void retry(IORunnable runnable, long retryIntervalMillis, long retryTimeoutMillis) throws IOException { - retry(IOSupplier.from(runnable), retryIntervalMillis, retryTimeoutMillis); - } - - private T retry(IOSupplier supplier, long retryIntervalMillis, long retryTimeoutMillis) throws IOException { - long start = System.nanoTime(); - long retryTimeoutNanos = TimeUnit.MILLISECONDS.toNanos(retryTimeoutMillis); - while (true) { - try { - return supplier.get(); - } catch (IOException e) { - if (System.nanoTime() - start > retryTimeoutNanos) { - failed = true; - LOGGER.error("Failed to execute IO operation, retry timeout", e); - throw e; - } - checkFailed(); - LOGGER.warn("Failed to execute IO operation, retrying in {}ms, error: {}", retryIntervalMillis, e.getMessage()); - Threads.sleep(retryIntervalMillis); - } - } - } - - private void checkFailed() throws IOException { - if (failed) { - IOException e = new IOException("Failed to execute IO operation, WAL failed"); - LOGGER.error("Failed to execute IO operation, WAL failed", e); - throw e; - } - } - - protected abstract void doWrite(ByteBuf src, long position) throws IOException; - - protected abstract void doFlush() throws IOException; - - protected abstract int doRead(ByteBuf dst, long position, int length) throws IOException; - - private interface IOSupplier { - T get() throws IOException; - - static IOSupplier from(IORunnable runnable) { - return () -> { - runnable.run(); - return null; - }; - } - } - - private interface IORunnable { - void run() throws IOException; - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALBlockDeviceChannel.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALBlockDeviceChannel.java deleted file mode 100644 index a625875886..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALBlockDeviceChannel.java +++ /dev/null @@ -1,339 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.s3.wal.exception.WALCapacityMismatchException; -import com.automq.stream.s3.wal.exception.WALNotInitializedException; -import com.automq.stream.thirdparty.moe.cnkirito.kdio.DirectIOLib; -import com.automq.stream.thirdparty.moe.cnkirito.kdio.DirectIOUtils; -import com.automq.stream.thirdparty.moe.cnkirito.kdio.DirectRandomAccessFile; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.File; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.util.concurrent.ExecutionException; - -import io.netty.buffer.ByteBuf; -import io.netty.util.concurrent.FastThreadLocal; - -import static com.automq.stream.s3.Constants.CAPACITY_NOT_SET; -import static com.automq.stream.s3.wal.util.WALUtil.isBlockDevice; - -public class WALBlockDeviceChannel extends AbstractWALChannel { - private static final Logger LOGGER = LoggerFactory.getLogger(WALBlockDeviceChannel.class); - private static final String CHECK_DIRECT_IO_AVAILABLE_FORMAT = "%s.check_direct_io_available"; - final String path; - final long capacityWant; - final boolean recoveryMode; - final DirectIOLib directIOLib; - /** - * 0 means allocate on demand - */ - final int initTempBufferSize; - /** - * 0 means no limit - */ - final int maxTempBufferSize; - /** - * Flag indicating whether unaligned write is allowed. - * Currently, it is only allowed when testing. - */ - public boolean unalignedWrite = false; - - long capacityFact = 0; - DirectRandomAccessFile randomAccessFile; - - FastThreadLocal threadLocalByteBuffer = new FastThreadLocal<>() { - @Override - protected ByteBuffer initialValue() { - return DirectIOUtils.allocateForDirectIO(directIOLib, initTempBufferSize); - } - }; - - public WALBlockDeviceChannel(String path, long capacityWant) { - this(path, capacityWant, 0, 0, false); - } - - public WALBlockDeviceChannel(String path, long capacityWant, int initTempBufferSize, int maxTempBufferSize, - boolean recoveryMode) { - this.path = path; - this.recoveryMode = recoveryMode; - if (recoveryMode) { - this.capacityWant = CAPACITY_NOT_SET; - } else { - assert capacityWant > 0; - this.capacityWant = capacityWant; - if (!WALUtil.isAligned(capacityWant)) { - throw new RuntimeException("wal capacity must be aligned by block size when using block device"); - } - } - this.initTempBufferSize = initTempBufferSize; - this.maxTempBufferSize = maxTempBufferSize; - - DirectIOLib lib = DirectIOLib.getLibForPath(path); - if (null == lib) { - throw new RuntimeException("O_DIRECT not supported"); - } - int blockSize = lib.blockSize(); - if (WALUtil.BLOCK_SIZE % blockSize != 0) { - throw new RuntimeException(String.format("block size %d is not a multiple of %d, update it by jvm option: -D%s=%d", - WALUtil.BLOCK_SIZE, blockSize, WALUtil.BLOCK_SIZE_PROPERTY, blockSize)); - } - this.directIOLib = lib; - } - - /** - * Check whether the {@link WALBlockDeviceChannel} is available for the given path. - * - * @return null if available, otherwise the reason why it's not available - */ - public static String checkAvailable(String path) { - if (!DirectIOLib.binit) { - return "O_DIRECT not supported"; - } - if (!DirectIOUtils.allocatorAvailable()) { - return "java.nio.DirectByteBuffer.(long, int) not available." + - " Add --add-opens=java.base/java.nio=ALL-UNNAMED and -Dio.netty.tryReflectionSetAccessible=true to JVM options may fix this."; - } - if (!isBlockDevice(path)) { - String reason = tryOpenFileWithDirectIO(String.format(CHECK_DIRECT_IO_AVAILABLE_FORMAT, path)); - if (null != reason) { - return "O_DIRECT not supported by the file system, path: " + path + ", reason: " + reason; - } - } - return null; - } - - /** - * Try to open a file with O_DIRECT flag to check whether the file system supports O_DIRECT. - * The file will be deleted after the test. - * - * @return null if the file is opened successfully, otherwise the reason why it's not available - */ - private static String tryOpenFileWithDirectIO(String path) { - File file = new File(path); - try { - DirectRandomAccessFile randomAccessFile = new DirectRandomAccessFile(file, "rw"); - randomAccessFile.close(); - return null; - } catch (IOException e) { - return e.getMessage(); - } finally { - // the file may be created in {@link DirectRandomAccessFile(File, String)}, so delete it - file.delete(); - } - } - - @Override - public void open(CapacityReader reader) throws IOException { - if (!isBlockDevice(path)) { - openAndCheckFile(); - } else { - try { - long capacity = WALUtil.getBlockDeviceCapacity(path); - if (!recoveryMode && capacityWant > capacity) { - // the real capacity of the block device is smaller than requested - throw new WALCapacityMismatchException(path, capacityWant, capacity); - } - } catch (ExecutionException e) { - LOGGER.warn("failed to get the real capacity of the block device {}, just skip checking", path, e); - } - // We could not get the real capacity of the WAL in block device, so we just use the `capacityWant` as the capacity here - // It will be checked and updated in `checkCapacity` later - capacityFact = capacityWant; - } - - randomAccessFile = new DirectRandomAccessFile(new File(path), "rw"); - - checkCapacity(reader); - } - - /** - * Create the file and set length if not exists, and check the file size if exists. - */ - private void openAndCheckFile() throws IOException { - File file = new File(path); - if (file.exists()) { - if (!file.isFile()) { - throw new IOException(path + " is not a file"); - } - capacityFact = file.length(); - if (!recoveryMode && capacityFact != capacityWant) { - // the file exists but not the same size as requested - throw new WALCapacityMismatchException(path, capacityWant, capacityFact); - } - } else { - // the file does not exist - if (recoveryMode) { - throw new WALNotInitializedException("try to open an uninitialized WAL in recovery mode: file not exists. path: " + path); - } - WALUtil.createFile(path, capacityWant); - capacityFact = capacityWant; - } - } - - private void checkCapacity(CapacityReader reader) throws IOException { - if (null == reader) { - return; - } - Long capacity = reader.capacity(this); - if (null == capacity) { - if (recoveryMode) { - throw new WALNotInitializedException("try to open an uninitialized WAL in recovery mode: empty header. path: " + path); - } - } else if (capacityFact == CAPACITY_NOT_SET) { - // recovery mode on block device - capacityFact = capacity; - } else if (capacityFact != capacity) { - throw new WALCapacityMismatchException(path, capacityFact, capacity); - } - assert capacityFact != CAPACITY_NOT_SET; - } - - @Override - public void close() { - try { - if (randomAccessFile != null) { - randomAccessFile.close(); - } - } catch (IOException ignored) { - } - } - - @Override - public long capacity() { - return capacityFact; - } - - @Override - public String path() { - return path; - } - - private ByteBuffer getBuffer(int alignedSize) { - assert WALUtil.isAligned(alignedSize); - - ByteBuffer currentBuf = threadLocalByteBuffer.get(); - if (alignedSize <= currentBuf.capacity()) { - return currentBuf; - } - if (maxTempBufferSize > 0 && alignedSize > maxTempBufferSize) { - throw new RuntimeException("too large write size"); - } - - ByteBuffer newBuf = DirectIOUtils.allocateForDirectIO(directIOLib, alignedSize); - threadLocalByteBuffer.set(newBuf); - DirectIOUtils.releaseDirectBuffer(currentBuf); - return newBuf; - } - - @Override - public void doWrite(ByteBuf src, long position) throws IOException { - if (unalignedWrite) { - // unaligned write, just used for testing - unalignedWrite(src, position); - return; - } - assert WALUtil.isAligned(position); - - int alignedSize = (int) WALUtil.alignLargeByBlockSize(src.readableBytes()); - assert position + alignedSize <= capacity(); - ByteBuffer tmpBuf = getBuffer(alignedSize); - tmpBuf.clear(); - - for (ByteBuffer buffer : src.nioBuffers()) { - tmpBuf.put(buffer); - } - tmpBuf.position(0).limit(alignedSize); - - write(tmpBuf, position); - } - - private void unalignedWrite(ByteBuf src, long position) throws IOException { - long start = position; - long end = position + src.readableBytes(); - long alignedStart = WALUtil.alignSmallByBlockSize(start); - long alignedEnd = WALUtil.alignLargeByBlockSize(end); - int alignedSize = (int) (alignedEnd - alignedStart); - - // read the data in the range [alignedStart, alignedEnd) to tmpBuf - ByteBuffer tmpBuf = getBuffer(alignedSize); - tmpBuf.position(0).limit(alignedSize); - read(tmpBuf, alignedStart); - - // overwrite the data in the range [start, end) in tmpBuf - for (ByteBuffer buffer : src.nioBuffers()) { - tmpBuf.position((int) (start - alignedStart)); - start += buffer.remaining(); - tmpBuf.put(buffer); - } - tmpBuf.position(0).limit(alignedSize); - - // write it - write(tmpBuf, alignedStart); - } - - private int write(ByteBuffer src, long position) throws IOException { - assert WALUtil.isAligned(src.remaining()); - - int bytesWritten = 0; - while (src.hasRemaining()) { - int written = randomAccessFile.write(src, position + bytesWritten); - // kdio will throw an exception rather than return -1, so we don't need to check for -1 - bytesWritten += written; - } - return bytesWritten; - } - - @Override - public void doFlush() { - } - - @Override - public int doRead(ByteBuf dst, long position, int length) throws IOException { - long start = position; - length = Math.min(length, dst.writableBytes()); - long end = position + length; - long alignedStart = WALUtil.alignSmallByBlockSize(start); - long alignedEnd = WALUtil.alignLargeByBlockSize(end); - int alignedSize = (int) (alignedEnd - alignedStart); - // capacity may be CAPACITY_NOT_SET only when we call {@link CapacityReader#capacity} in recovery mode - assert CAPACITY_NOT_SET == capacity() || alignedEnd <= capacity(); - - ByteBuffer tmpBuf = getBuffer(alignedSize); - tmpBuf.position(0).limit(alignedSize); - - read(tmpBuf, alignedStart); - tmpBuf.position((int) (start - alignedStart)).limit((int) (end - alignedStart)); - - dst.writeBytes(tmpBuf); - return (int) (end - start); - } - - private int read(ByteBuffer dst, long position) throws IOException { - int bytesRead = 0; - while (dst.hasRemaining()) { - int read = randomAccessFile.read(dst, position + bytesRead); - // kdio will throw an exception rather than return -1, so we don't need to check for -1 - bytesRead += read; - } - return bytesRead; - } - - @Override - public boolean useDirectIO() { - return true; - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALCachedChannel.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALCachedChannel.java deleted file mode 100644 index 1d91766fa0..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALCachedChannel.java +++ /dev/null @@ -1,167 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.s3.ByteBufAlloc; - -import java.io.IOException; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.Constants.CAPACITY_NOT_SET; - -/** - * A wrapper of {@link WALChannel} that caches for read to reduce I/O. - */ -public class WALCachedChannel implements WALChannel { - - private static final int DEFAULT_CACHE_SIZE = 1 << 20; - - private final WALChannel channel; - private final int cacheSize; - - private ByteBuf cache; - private long cachePosition = -1; - - private WALCachedChannel(WALChannel channel, int cacheSize) { - this.channel = channel; - this.cacheSize = cacheSize; - } - - public static WALCachedChannel of(WALChannel channel) { - return new WALCachedChannel(channel, DEFAULT_CACHE_SIZE); - } - - public static WALCachedChannel of(WALChannel channel, int cacheSize) { - return new WALCachedChannel(channel, cacheSize); - } - - @Override - public int read(ByteBuf dst, long position, int length) throws IOException { - return read(channel::read, dst, position, length); - } - - @Override - public int retryRead(ByteBuf dst, long position, int length, long retryIntervalMillis, - long retryTimeoutMillis) throws IOException { - Reader reader = (buf, pos, len) -> channel.retryRead(buf, pos, len, retryIntervalMillis, retryTimeoutMillis); - return read(reader, dst, position, length); - } - - /** - * As we use a common cache for all threads, we need to synchronize the read. - */ - private synchronized int read(Reader reader, ByteBuf dst, long position, int length) throws IOException { - if (CAPACITY_NOT_SET == channel.capacity()) { - // If we don't know the capacity now, we can't cache. - return reader.read(dst, position, length); - } - - long start = position; - length = Math.min(length, dst.writableBytes()); - long end = position + length; - - ByteBuf cache = getCache(); - if (length > cache.capacity()) { - // If the length is larger than the cache capacity, we can't cache. - return reader.read(dst, position, length); - } - - boolean fallWithinCache = cachePosition >= 0 && cachePosition <= start && end <= cachePosition + cache.readableBytes(); - if (!fallWithinCache) { - cache.clear(); - cachePosition = start; - // Make sure the cache is not larger than the channel capacity. - int cacheLength = (int) Math.min(cache.writableBytes(), channel.capacity() - cachePosition); - reader.read(cache, cachePosition, cacheLength); - } - - // Now the cache is ready. - int relativePosition = (int) (start - cachePosition); - dst.writeBytes(cache, relativePosition, length); - return length; - } - - @Override - public void close() { - releaseCache(); - this.channel.close(); - } - - /** - * Release the cache if it is not null. - * This method should be called when no more {@link #read}s will be called to release the allocated memory. - */ - public synchronized void releaseCache() { - if (this.cache != null) { - this.cache.release(); - this.cache = null; - } - this.cachePosition = -1; - } - - /** - * Get the cache. If the cache is not initialized, initialize it. - * Should be called under synchronized. - */ - private ByteBuf getCache() { - if (this.cache == null) { - this.cache = ByteBufAlloc.byteBuffer(cacheSize); - } - return this.cache; - } - - private interface Reader { - int read(ByteBuf dst, long position, int length) throws IOException; - } - - @Override - public void open(CapacityReader reader) throws IOException { - this.channel.open(reader); - } - - @Override - public long capacity() { - return this.channel.capacity(); - } - - @Override - public String path() { - return this.channel.path(); - } - - @Override - public void write(ByteBuf src, long position) throws IOException { - this.channel.write(src, position); - } - - @Override - public void retryWrite(ByteBuf src, long position, long retryIntervalMillis, - long retryTimeoutMillis) throws IOException { - channel.retryWrite(src, position, retryIntervalMillis, retryTimeoutMillis); - } - - @Override - public void flush() throws IOException { - this.channel.flush(); - } - - @Override - public void retryFlush(long retryIntervalMillis, long retryTimeoutMillis) throws IOException { - channel.retryFlush(retryIntervalMillis, retryTimeoutMillis); - } - - @Override - public boolean useDirectIO() { - return channel.useDirectIO(); - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALChannel.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALChannel.java deleted file mode 100644 index b5259b2e4a..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALChannel.java +++ /dev/null @@ -1,200 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.s3.wal.exception.WALCapacityMismatchException; -import com.automq.stream.s3.wal.exception.WALNotInitializedException; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.File; -import java.io.IOException; -import java.util.concurrent.TimeUnit; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.Constants.CAPACITY_NOT_SET; -import static com.automq.stream.s3.wal.util.WALUtil.isBlockDevice; - -/** - * There are two implementations of WALChannel: - * 1. WALFileChannel based on file system, which calls fsync after each write to ensure data is flushed to disk. - * 2. WALBlockDeviceChannel based on block device, which uses O_DIRECT to bypass page cache. - */ -public interface WALChannel { - - long DEFAULT_RETRY_INTERVAL = TimeUnit.MILLISECONDS.toMillis(100); - long DEFAULT_RETRY_TIMEOUT = TimeUnit.MINUTES.toMillis(1); - - static WALChannelBuilder builder(String path) { - return new WALChannelBuilder(path); - } - - /** - * Open the channel for read and write. - * If {@code reader} is null, checks will be skipped. - * - * @param reader the reader to get the capacity of the channel - * @throws WALCapacityMismatchException if the capacity of the channel does not match the expected capacity - * @throws WALNotInitializedException if try to open an un-initialized channel in recovery mode - * @throws IOException if any I/O error happens - */ - void open(CapacityReader reader) throws IOException; - - default void open() throws IOException { - open(null); - } - - void close(); - - long capacity(); - - String path(); - - /** - * Write bytes from the given buffer to the given position of the channel from the current reader index - * to the end of the buffer. It only returns when all bytes are written successfully. - * {@link #flush()} should be called after this method to ensure data is flushed to disk. - * This method will change the reader index of the given buffer to the end of the written bytes. - * This method will not change the writer index of the given buffer. - */ - void write(ByteBuf src, long position) throws IOException; - - default void retryWrite(ByteBuf src, long position) throws IOException { - retryWrite(src, position, DEFAULT_RETRY_INTERVAL, DEFAULT_RETRY_TIMEOUT); - } - - /** - * Retry {@link #write(ByteBuf, long)} with the given interval until success or timeout. - */ - void retryWrite(ByteBuf src, long position, long retryIntervalMillis, long retryTimeoutMillis) throws IOException; - - /** - * Flush to disk. - */ - void flush() throws IOException; - - default void retryFlush() throws IOException { - retryFlush(DEFAULT_RETRY_INTERVAL, DEFAULT_RETRY_TIMEOUT); - } - - /** - * Retry {@link #flush()} with the given interval until success or timeout. - */ - void retryFlush(long retryIntervalMillis, long retryTimeoutMillis) throws IOException; - - default int read(ByteBuf dst, long position) throws IOException { - return read(dst, position, dst.writableBytes()); - } - - /** - * Read bytes from the given position of the channel to the given buffer from the current writer index - * until reaching the given length or the end of the channel. - * This method will change the writer index of the given buffer to the end of the read bytes. - * This method will not change the reader index of the given buffer. - * If the given length is larger than the writable bytes of the given buffer, only the first - * {@code dst.writableBytes()} bytes will be read. - */ - int read(ByteBuf dst, long position, int length) throws IOException; - - default int retryRead(ByteBuf dst, long position) throws IOException { - return retryRead(dst, position, dst.writableBytes(), DEFAULT_RETRY_INTERVAL, DEFAULT_RETRY_TIMEOUT); - } - - /** - * Retry {@link #read(ByteBuf, long, int)} with the given interval until success or timeout. - */ - int retryRead(ByteBuf dst, long position, int length, long retryIntervalMillis, long retryTimeoutMillis) throws IOException; - - boolean useDirectIO(); - - interface CapacityReader { - /** - * Get the capacity of the given channel. - * It returns null if the channel has not been initialized before. - */ - Long capacity(WALChannel channel) throws IOException; - } - - class WALChannelBuilder { - private static final Logger LOGGER = LoggerFactory.getLogger(WALChannelBuilder.class); - private final String path; - private Boolean direct; - private long capacity; - private int initBufferSize; - private int maxBufferSize; - private boolean recoveryMode; - - private WALChannelBuilder(String path) { - this.path = path; - } - - public WALChannelBuilder direct(boolean direct) { - this.direct = direct; - return this; - } - - public WALChannelBuilder capacity(long capacity) { - assert capacity == CAPACITY_NOT_SET || WALUtil.isAligned(capacity); - this.capacity = capacity; - return this; - } - - public WALChannelBuilder initBufferSize(int initBufferSize) { - this.initBufferSize = initBufferSize; - return this; - } - - public WALChannelBuilder maxBufferSize(int maxBufferSize) { - this.maxBufferSize = maxBufferSize; - return this; - } - - public WALChannelBuilder recoveryMode(boolean recoveryMode) { - this.recoveryMode = recoveryMode; - return this; - } - - public WALChannel build() { - String directNotAvailableMsg = WALBlockDeviceChannel.checkAvailable(path); - boolean isBlockDevice = isBlockDevice(path); - boolean useDirect = false; - if (direct != null) { - // Set by user. - useDirect = direct; - } else if (isBlockDevice) { - // We can only use direct IO for block devices. - useDirect = true; - } else if (directNotAvailableMsg == null) { - // If direct IO is available, we use it by default. - useDirect = true; - } - - if (useDirect && directNotAvailableMsg != null) { - throw new IllegalArgumentException(directNotAvailableMsg); - } - - if (!isBlockDevice) { - LOGGER.warn("WAL in a file system, which may cause performance degradation. path: {}", new File(path).getAbsolutePath()); - } - - if (useDirect) { - return new WALBlockDeviceChannel(path, capacity, initBufferSize, maxBufferSize, recoveryMode); - } else { - LOGGER.warn("Direct IO not used for WAL, which may cause performance degradation. path: {}, isBlockDevice: {}, reason: {}", - new File(path).getAbsolutePath(), isBlockDevice, directNotAvailableMsg); - return new WALFileChannel(path, capacity, recoveryMode); - } - } - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALFileChannel.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALFileChannel.java deleted file mode 100644 index 4057fbecdb..0000000000 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALFileChannel.java +++ /dev/null @@ -1,158 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.s3.wal.exception.WALCapacityMismatchException; -import com.automq.stream.s3.wal.exception.WALNotInitializedException; - -import java.io.File; -import java.io.IOException; -import java.io.RandomAccessFile; -import java.nio.ByteBuffer; -import java.nio.channels.FileChannel; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.Constants.CAPACITY_NOT_SET; - -public class WALFileChannel extends AbstractWALChannel { - final String filePath; - final long fileCapacityWant; - /** - * When set to true, the file should exist and the file size does not need to be verified. - */ - final boolean recoveryMode; - long fileCapacityFact = 0; - RandomAccessFile randomAccessFile; - FileChannel fileChannel; - - public WALFileChannel(String filePath, long fileCapacityWant, boolean recoveryMode) { - this.filePath = filePath; - this.recoveryMode = recoveryMode; - if (recoveryMode) { - this.fileCapacityWant = CAPACITY_NOT_SET; - } else { - assert fileCapacityWant > 0; - this.fileCapacityWant = fileCapacityWant; - } - } - - @Override - public void open(CapacityReader reader) throws IOException { - File file = new File(filePath); - if (file.exists()) { - if (!file.isFile()) { - throw new IOException(filePath + " is not a file"); - } - randomAccessFile = new RandomAccessFile(file, "rw"); - fileCapacityFact = randomAccessFile.length(); - if (!recoveryMode && fileCapacityFact != fileCapacityWant) { - // the file exists but not the same size as requested - throw new WALCapacityMismatchException(filePath, fileCapacityWant, fileCapacityFact); - } - } else { - // the file does not exist - if (recoveryMode) { - throw new WALNotInitializedException("try to open an uninitialized WAL in recovery mode: file not exists: " + filePath); - } - WALUtil.createFile(filePath, fileCapacityWant); - randomAccessFile = new RandomAccessFile(filePath, "rw"); - fileCapacityFact = fileCapacityWant; - } - - fileChannel = randomAccessFile.getChannel(); - - checkCapacity(reader); - } - - private void checkCapacity(CapacityReader reader) throws IOException { - if (null == reader) { - return; - } - Long capacity = reader.capacity(this); - if (null == capacity) { - if (recoveryMode) { - throw new WALNotInitializedException("try to open an uninitialized WAL in recovery mode: empty header. path: " + filePath); - } - } else if (fileCapacityFact != capacity) { - throw new WALCapacityMismatchException(filePath, fileCapacityFact, capacity); - } - assert fileCapacityFact != CAPACITY_NOT_SET; - } - - @Override - public void close() { - try { - fileChannel.close(); - randomAccessFile.close(); - } catch (IOException ignored) { - } - } - - @Override - public long capacity() { - return fileCapacityFact; - } - - @Override - public String path() { - return filePath; - } - - @Override - public void doWrite(ByteBuf src, long position) throws IOException { - assert src.readableBytes() + position <= capacity(); - ByteBuffer[] nioBuffers = src.nioBuffers(); - for (ByteBuffer nioBuffer : nioBuffers) { - int bytesWritten = write(nioBuffer, position); - position += bytesWritten; - } - } - - @Override - public void doFlush() throws IOException { - fileChannel.force(false); - } - - @Override - public int doRead(ByteBuf dst, long position, int length) throws IOException { - length = Math.min(length, dst.writableBytes()); - assert position + length <= capacity(); - int bytesRead = 0; - while (dst.isWritable()) { - int read = dst.writeBytes(fileChannel, position + bytesRead, length); - if (read == -1) { - // EOF - break; - } - bytesRead += read; - } - return bytesRead; - } - - private int write(ByteBuffer src, long position) throws IOException { - int bytesWritten = 0; - while (src.hasRemaining()) { - int written = fileChannel.write(src, position + bytesWritten); - if (written == -1) { - throw new IOException("write -1"); - } - bytesWritten += written; - } - return bytesWritten; - } - - @Override - public boolean useDirectIO() { - return false; - } -} diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java index ab0ceacaca..2be513bb10 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/util/WALUtil.java @@ -1,21 +1,15 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - package com.automq.stream.s3.wal.util; +import com.automq.stream.s3.ByteBufAlloc; import com.automq.stream.s3.wal.common.Record; import com.automq.stream.s3.wal.common.RecordHeader; +import com.automq.stream.s3.wal.exception.WALCapacityMismatchException; import com.automq.stream.utils.CommandResult; import com.automq.stream.utils.CommandUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import java.io.File; import java.io.IOException; import java.io.RandomAccessFile; @@ -24,9 +18,9 @@ import java.util.zip.CRC32; import io.netty.buffer.ByteBuf; +import jnr.posix.POSIX; import jnr.posix.POSIXFactory; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_MAGIC_CODE; import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; public class WALUtil { @@ -36,21 +30,31 @@ public class WALUtil { "4096" )); + private static final Logger LOGGER = LoggerFactory.getLogger(WALUtil.class); + public static Record generateRecord(ByteBuf body, ByteBuf emptyHeader, int crc, long start) { - return generateRecord(body, emptyHeader, crc, start, true); + crc = 0 == crc ? WALUtil.crc32(body) : crc; + ByteBuf header = new RecordHeader(start, body.readableBytes(), crc).marshal(emptyHeader); + return new Record(header, body); } - public static Record generateRecord(ByteBuf body, ByteBuf emptyHeader, int crc, long start, boolean calculateCRC) { - crc = 0 == crc ? WALUtil.crc32(body) : crc; - ByteBuf header = new RecordHeader() - .setMagicCode(RECORD_HEADER_MAGIC_CODE) - .setRecordBodyLength(body.readableBytes()) - .setRecordBodyOffset(start + RECORD_HEADER_SIZE) - .setRecordBodyCRC(crc) - .marshal(emptyHeader, calculateCRC); + public static Record generatePaddingRecord(ByteBuf emptyHeader, long start, int length) { + int bodyLength = length - RECORD_HEADER_SIZE; + + ByteBuf header = new RecordHeader(start, bodyLength).marshal(emptyHeader); + + ByteBuf body = ByteBufAlloc.byteBuffer(bodyLength); + body.writeZero(bodyLength); + return new Record(header, body); } + public static ByteBuf generateHeader(ByteBuf body, ByteBuf emptyHeader, int crc, long start) { + crc = 0 == crc ? WALUtil.crc32(body) : crc; + return new RecordHeader(start, body.readableBytes(), crc) + .marshal(emptyHeader); + } + /** * Get CRC32 of the given ByteBuf from current reader index to the end. * This method will not change the reader index of the given ByteBuf. @@ -102,11 +106,15 @@ public static boolean isAligned(long offset) { return offset % BLOCK_SIZE == 0; } + public static RandomAccessFile createFile(String path, long length) throws IOException { + return createFile(path, length, "rw"); + } + /** * Create a file with the given path and length. * Note {@code path} must NOT exist. */ - public static void createFile(String path, long length) throws IOException { + public static RandomAccessFile createFile(String path, long length, String openMode) throws IOException { File file = new File(path); assert !file.exists(); @@ -114,20 +122,16 @@ public static void createFile(String path, long length) throws IOException { if (null != parent && !parent.exists() && !parent.mkdirs()) { throw new IOException("mkdirs " + parent + " fail"); } - if (!file.createNewFile()) { - throw new IOException("create " + path + " fail"); - } - if (!file.setReadable(true)) { - throw new IOException("set " + path + " readable fail"); - } - if (!file.setWritable(true)) { - throw new IOException("set " + path + " writable fail"); - } - - // set length - try (RandomAccessFile raf = new RandomAccessFile(file, "rw")) { + RandomAccessFile raf = new RandomAccessFile(file, openMode); + long realLength = raf.length(); + if (realLength == 0) { + // set length raf.setLength(length); + } else if (realLength != length) { + // the file exists but not the same size as requested + throw new WALCapacityMismatchException(path, length, realLength); } + return raf; } /** @@ -158,15 +162,25 @@ public static boolean isBlockDevice(String path) { if (!new File(path).exists()) { return false; } - boolean isBlockDevice; + + POSIX posix; + try { + posix = POSIXFactory.getNativePOSIX(); + } catch (Exception e) { + LOGGER.warn("Failed to get native POSIX, fallback to check by prefix", e); + return isBlockDeviceByPrefix(path); + } + try { - isBlockDevice = POSIXFactory.getPOSIX() - .stat(path) - .isBlockDev(); + return posix.stat(path).isBlockDev(); } catch (Exception e) { // In some OS (like Windows), the isBlockDev() method may throw an IllegalStateException. - isBlockDevice = false; + LOGGER.warn("Failed to check if {} is a block device, fallback to check by prefix", path, e); + return isBlockDeviceByPrefix(path); } - return isBlockDevice; + } + + private static boolean isBlockDeviceByPrefix(String path) { + return path.startsWith("/dev/"); } } diff --git a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannel.java b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannel.java index 49abdb9e4d..9e61499b08 100755 --- a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannel.java +++ b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannel.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.thirdparty.moe.cnkirito.kdio; diff --git a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannelImpl.java b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannelImpl.java index b0e6bcd5f8..27cb7a935e 100755 --- a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannelImpl.java +++ b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectChannelImpl.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.thirdparty.moe.cnkirito.kdio; diff --git a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOLib.java b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOLib.java index 196a17b05a..dd39724738 100755 --- a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOLib.java +++ b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOLib.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.thirdparty.moe.cnkirito.kdio; diff --git a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOUtils.java b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOUtils.java index 2beaa45e5b..d18cb3b7c8 100755 --- a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOUtils.java +++ b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectIOUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.thirdparty.moe.cnkirito.kdio; diff --git a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectRandomAccessFile.java b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectRandomAccessFile.java index 589f5ec016..cf42739137 100755 --- a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectRandomAccessFile.java +++ b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/DirectRandomAccessFile.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.thirdparty.moe.cnkirito.kdio; @@ -46,6 +54,10 @@ public DirectRandomAccessFile(File file, String mode) this.channel = DirectChannelImpl.getChannel(file, readOnly); } + public DirectRandomAccessFile(DirectChannel channel) { + this.channel = channel; + } + @Override public void close() throws IOException { channel.close(); diff --git a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/OpenFlags.java b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/OpenFlags.java index 48a0dd0177..e9f955b5b5 100755 --- a/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/OpenFlags.java +++ b/s3stream/src/main/java/com/automq/stream/thirdparty/moe/cnkirito/kdio/OpenFlags.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.thirdparty.moe.cnkirito.kdio; diff --git a/s3stream/src/main/java/com/automq/stream/utils/Arguments.java b/s3stream/src/main/java/com/automq/stream/utils/Arguments.java index 94a6bf91dc..f137fbea88 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/Arguments.java +++ b/s3stream/src/main/java/com/automq/stream/utils/Arguments.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/AsyncRateLimiter.java b/s3stream/src/main/java/com/automq/stream/utils/AsyncRateLimiter.java index 80a75e2011..0140e58f47 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/AsyncRateLimiter.java +++ b/s3stream/src/main/java/com/automq/stream/utils/AsyncRateLimiter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/AsyncSemaphore.java b/s3stream/src/main/java/com/automq/stream/utils/AsyncSemaphore.java similarity index 61% rename from s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/AsyncSemaphore.java rename to s3stream/src/main/java/com/automq/stream/utils/AsyncSemaphore.java index 9b7936546b..3a0df0e068 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/cache/blockcache/AsyncSemaphore.java +++ b/s3stream/src/main/java/com/automq/stream/utils/AsyncSemaphore.java @@ -1,17 +1,23 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ -package com.automq.stream.s3.cache.blockcache; - -import com.automq.stream.utils.threads.EventLoop; +package com.automq.stream.utils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -19,9 +25,10 @@ import java.util.LinkedList; import java.util.Queue; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; import java.util.function.Supplier; -class AsyncSemaphore { +public class AsyncSemaphore { private static final Logger LOGGER = LoggerFactory.getLogger(AsyncSemaphore.class); private final Queue tasks = new LinkedList<>(); private long permits; @@ -36,11 +43,11 @@ public AsyncSemaphore(long permits) { * @param requiredPermits the required permits * @param task task to run when the permits are available, the task should return a CompletableFuture * which will be completed when the permits could be released. - * @param eventLoop the eventLoop to run the task when the permits are available + * @param executor the executor to run the task when the permits are available * @return true if the permits are acquired, false if the task is added to the waiting queue. */ public synchronized boolean acquire(long requiredPermits, Supplier> task, - EventLoop eventLoop) { + Executor executor) { if (permits >= 0) { // allow permits minus to negative permits -= requiredPermits; @@ -48,10 +55,11 @@ public synchronized boolean acquire(long requiredPermits, Supplier release(requiredPermits)); } catch (Throwable e) { LOGGER.error("Error in task", e); + release(requiredPermits); } return true; } else { - tasks.add(new AsyncSemaphoreTask(requiredPermits, task, eventLoop)); + tasks.add(new AsyncSemaphoreTask(requiredPermits, task, executor)); return false; } } @@ -69,8 +77,8 @@ synchronized void release(long requiredPermits) { if (permits > 0) { AsyncSemaphoreTask t = tasks.poll(); if (t != null) { - // use eventLoop to reset the thread stack to avoid stack overflow - t.eventLoop.execute(() -> acquire(t.requiredPermits, t.task, t.eventLoop)); + // use executor to reset the thread stack to avoid stack overflow + t.executor.execute(() -> acquire(t.requiredPermits, t.task, t.executor)); } } } @@ -78,12 +86,12 @@ synchronized void release(long requiredPermits) { static class AsyncSemaphoreTask { final long requiredPermits; final Supplier> task; - final EventLoop eventLoop; + final Executor executor; - public AsyncSemaphoreTask(long requiredPermits, Supplier> task, EventLoop eventLoop) { + public AsyncSemaphoreTask(long requiredPermits, Supplier> task, Executor executor) { this.requiredPermits = requiredPermits; this.task = task; - this.eventLoop = eventLoop; + this.executor = executor; } } } diff --git a/s3stream/src/main/java/com/automq/stream/utils/ByteBufInputStream.java b/s3stream/src/main/java/com/automq/stream/utils/ByteBufInputStream.java index 7d18cfe572..95f2596478 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/ByteBufInputStream.java +++ b/s3stream/src/main/java/com/automq/stream/utils/ByteBufInputStream.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/CloseableIterator.java b/s3stream/src/main/java/com/automq/stream/utils/CloseableIterator.java index 97b027674d..f0c3250050 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/CloseableIterator.java +++ b/s3stream/src/main/java/com/automq/stream/utils/CloseableIterator.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/CollectionHelper.java b/s3stream/src/main/java/com/automq/stream/utils/CollectionHelper.java index ace5a3bb17..ef5ff82114 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/CollectionHelper.java +++ b/s3stream/src/main/java/com/automq/stream/utils/CollectionHelper.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/CommandResult.java b/s3stream/src/main/java/com/automq/stream/utils/CommandResult.java index 470277c007..6eecd1c6e7 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/CommandResult.java +++ b/s3stream/src/main/java/com/automq/stream/utils/CommandResult.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/CommandUtils.java b/s3stream/src/main/java/com/automq/stream/utils/CommandUtils.java index bf58da7e58..f90d2f7e07 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/CommandUtils.java +++ b/s3stream/src/main/java/com/automq/stream/utils/CommandUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; @@ -19,8 +27,12 @@ public class CommandUtils { public static CommandResult run(String... cmd) { + return run(runtime -> runtime.exec(cmd)); + } + + public static CommandResult run(CommandCall call) { try { - Process p = Runtime.getRuntime().exec(cmd); + Process p = call.call(Runtime.getRuntime()); try (BufferedReader inputReader = new BufferedReader(new InputStreamReader(p.getInputStream(), Charset.defaultCharset())); BufferedReader errorReader = new BufferedReader(new InputStreamReader(p.getErrorStream(), Charset.defaultCharset()))) { String stdout = inputReader.lines().collect(Collectors.joining("\n")); @@ -33,4 +45,8 @@ public static CommandResult run(String... cmd) { } } + public interface CommandCall { + Process call(Runtime runtime) throws IOException; + } + } diff --git a/s3stream/src/main/java/com/automq/stream/utils/ExceptionUtil.java b/s3stream/src/main/java/com/automq/stream/utils/ExceptionUtil.java new file mode 100644 index 0000000000..386d155e92 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/utils/ExceptionUtil.java @@ -0,0 +1,23 @@ +package com.automq.stream.utils; + +import java.util.Collection; +import java.util.Iterator; + +public class ExceptionUtil { + + /** + * Combine multiple exceptions into one, where the first exception is the primary one and the rest are suppressed. + * It returns null if the input collection is null or empty. + */ + public static E combine(Collection exceptions) { + if (null == exceptions || exceptions.isEmpty()) { + return null; + } + Iterator it = exceptions.iterator(); + E primary = it.next(); + while (it.hasNext()) { + primary.addSuppressed(it.next()); + } + return primary; + } +} diff --git a/s3stream/src/main/java/com/automq/stream/utils/FutureTicker.java b/s3stream/src/main/java/com/automq/stream/utils/FutureTicker.java index eec8cda15d..7069080ff8 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/FutureTicker.java +++ b/s3stream/src/main/java/com/automq/stream/utils/FutureTicker.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/FutureUtil.java b/s3stream/src/main/java/com/automq/stream/utils/FutureUtil.java index cc4f8c7e9b..d18685ab93 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/FutureUtil.java +++ b/s3stream/src/main/java/com/automq/stream/utils/FutureUtil.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; @@ -15,6 +23,7 @@ import org.slf4j.LoggerFactory; import java.util.Iterator; +import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutionException; @@ -23,6 +32,8 @@ import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.Stream; import io.netty.util.HashedWheelTimer; import io.netty.util.Timeout; @@ -157,4 +168,11 @@ public static CompletableFuture timeoutWithNewReturn(CompletableFuture return newCf; } + public static List> timeoutAndSilence(Stream> stream, long timeout, TimeUnit timeUnit) { + return stream.map(l -> timeoutAndSilence(l, timeout, timeUnit)).collect(Collectors.toList()); + } + + public static CompletableFuture timeoutAndSilence(CompletableFuture cf, long timeout, TimeUnit timeUnit) { + return cf.orTimeout(timeout, timeUnit).exceptionally(ex -> null); + } } diff --git a/s3stream/src/main/java/com/automq/stream/utils/GlobalSwitch.java b/s3stream/src/main/java/com/automq/stream/utils/GlobalSwitch.java index 9d5252596d..ec8f418f60 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/GlobalSwitch.java +++ b/s3stream/src/main/java/com/automq/stream/utils/GlobalSwitch.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/IdURI.java b/s3stream/src/main/java/com/automq/stream/utils/IdURI.java index 04b71d4c6c..a4f404c09c 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/IdURI.java +++ b/s3stream/src/main/java/com/automq/stream/utils/IdURI.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/LogContext.java b/s3stream/src/main/java/com/automq/stream/utils/LogContext.java index 8ef731c2d2..39dba29e13 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/LogContext.java +++ b/s3stream/src/main/java/com/automq/stream/utils/LogContext.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/LogSuppressor.java b/s3stream/src/main/java/com/automq/stream/utils/LogSuppressor.java index 93697f4bec..5a6e26b149 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/LogSuppressor.java +++ b/s3stream/src/main/java/com/automq/stream/utils/LogSuppressor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/SecretUtils.java b/s3stream/src/main/java/com/automq/stream/utils/SecretUtils.java index 1a0af6cb77..5bea54d6fc 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/SecretUtils.java +++ b/s3stream/src/main/java/com/automq/stream/utils/SecretUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/Systems.java b/s3stream/src/main/java/com/automq/stream/utils/Systems.java index 0d0be9a5e8..025196033e 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/Systems.java +++ b/s3stream/src/main/java/com/automq/stream/utils/Systems.java @@ -1,20 +1,32 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; import org.apache.commons.lang3.StringUtils; +import io.netty.util.internal.PlatformDependent; + public class Systems { public static final int CPU_CORES = Runtime.getRuntime().availableProcessors(); + public static final long HEAP_MEMORY_SIZE = Runtime.getRuntime().maxMemory(); + public static final long DIRECT_MEMORY_SIZE = PlatformDependent.maxDirectMemory(); public static long getEnvLong(String name, long defaultValue) { String value = System.getenv(name); diff --git a/s3stream/src/main/java/com/automq/stream/utils/ThreadUtils.java b/s3stream/src/main/java/com/automq/stream/utils/ThreadUtils.java index 449fc1d2e8..94353e349a 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/ThreadUtils.java +++ b/s3stream/src/main/java/com/automq/stream/utils/ThreadUtils.java @@ -1,19 +1,30 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; import org.slf4j.Logger; +import org.slf4j.helpers.NOPLogger; +import java.util.concurrent.ExecutorService; import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import io.netty.util.concurrent.FastThreadLocalThread; @@ -79,4 +90,33 @@ public static Runnable wrapRunnable(Runnable runnable, Logger logger) { } }; } + + /** + * A wrapper of {@link #shutdownExecutor} without logging. + */ + public static void shutdownExecutor(ExecutorService executorService, long timeout, TimeUnit timeUnit) { + shutdownExecutor(executorService, timeout, timeUnit, NOPLogger.NOP_LOGGER); + } + + /** + * Shuts down an executor service in two phases, first by calling shutdown to reject incoming tasks, + * and then calling shutdownNow, if necessary, to cancel any lingering tasks. + * After the timeout/on interrupt, the service is forcefully closed. + */ + public static void shutdownExecutor(ExecutorService executorService, long timeout, TimeUnit timeUnit, + Logger logger) { + if (null == executorService) { + return; + } + executorService.shutdown(); + try { + if (!executorService.awaitTermination(timeout, timeUnit)) { + executorService.shutdownNow(); + logger.error("Executor {} did not terminate in time, forcefully shutting down", executorService); + } + } catch (InterruptedException e) { + executorService.shutdownNow(); + Thread.currentThread().interrupt(); + } + } } diff --git a/s3stream/src/main/java/com/automq/stream/utils/Threads.java b/s3stream/src/main/java/com/automq/stream/utils/Threads.java index 591b994236..79ed77a3d2 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/Threads.java +++ b/s3stream/src/main/java/com/automq/stream/utils/Threads.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/ThrowableRunnable.java b/s3stream/src/main/java/com/automq/stream/utils/ThrowableRunnable.java index 2fad577d32..bb3ffe33f8 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/ThrowableRunnable.java +++ b/s3stream/src/main/java/com/automq/stream/utils/ThrowableRunnable.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/main/java/com/automq/stream/utils/URIUtils.java b/s3stream/src/main/java/com/automq/stream/utils/URIUtils.java index 1a0d9ab867..6e741b7713 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/URIUtils.java +++ b/s3stream/src/main/java/com/automq/stream/utils/URIUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; @@ -27,8 +35,8 @@ import java.util.regex.Pattern; public class URIUtils { - public static final Pattern URI_LIST_PATTERN = Pattern.compile("\\d+@.*?(?=,\\d+@|$)"); - public static final Pattern URI_PATTERN = Pattern.compile("(\\d+)@(.+)"); + public static final Pattern URI_LIST_PATTERN = Pattern.compile("[-]?\\d+@.*?(?=,[-]?\\d+@|$)"); + public static final Pattern URI_PATTERN = Pattern.compile("([-]?\\d+)@(.+)"); public static List parseIdURIList(String uriList) { if (StringUtils.isBlank(uriList)) { diff --git a/s3stream/src/main/java/com/automq/stream/utils/Utils.java b/s3stream/src/main/java/com/automq/stream/utils/Utils.java index eba9b6c4cf..8d9129b799 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/Utils.java +++ b/s3stream/src/main/java/com/automq/stream/utils/Utils.java @@ -1,16 +1,32 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; +import java.io.IOException; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.NoSuchFileException; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.attribute.BasicFileAttributes; + public class Utils { public static final String MAX_MERGE_READ_SPARSITY_RATE_NAME = "MERGE_READ_SPARSITY_RATE"; @@ -23,4 +39,41 @@ public static float getMaxMergeReadSparsityRate() { } return rate; } + + public static void delete(Path rootFile) throws IOException { + if (rootFile == null) + return; + Files.walkFileTree(rootFile, new SimpleFileVisitor() { + @Override + public FileVisitResult visitFileFailed(Path path, IOException exc) throws IOException { + if (exc instanceof NoSuchFileException) { + if (path.equals(rootFile)) { + // If the root path did not exist, ignore the error and terminate; + return FileVisitResult.TERMINATE; + } else { + // Otherwise, just continue walking as the file might already be deleted by other threads. + return FileVisitResult.CONTINUE; + } + } + throw exc; + } + + @Override + public FileVisitResult visitFile(Path path, BasicFileAttributes attrs) throws IOException { + Files.deleteIfExists(path); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path path, IOException exc) throws IOException { + // KAFKA-8999: if there's an exception thrown previously already, we should throw it + if (exc != null) { + throw exc; + } + + Files.deleteIfExists(path); + return FileVisitResult.CONTINUE; + } + }); + } } diff --git a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/AbstractOrderedCollection.java b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/AbstractOrderedCollection.java index 5b8f2d24a6..2c665dbe5c 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/AbstractOrderedCollection.java +++ b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/AbstractOrderedCollection.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.biniarysearch; diff --git a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/ComparableItem.java b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/ComparableItem.java index f4863506cd..71aacbe270 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/ComparableItem.java +++ b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/ComparableItem.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.biniarysearch; diff --git a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/IndexBlockOrderedBytes.java b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/IndexBlockOrderedBytes.java index 078ea2d29a..0cb4374494 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/IndexBlockOrderedBytes.java +++ b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/IndexBlockOrderedBytes.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.biniarysearch; diff --git a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamOffsetRangeSearchList.java b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamOffsetRangeSearchList.java index 44a13ad5d3..7062488dd7 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamOffsetRangeSearchList.java +++ b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamOffsetRangeSearchList.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.biniarysearch; diff --git a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamRecordBatchList.java b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamRecordBatchList.java index c09ea4b5d3..64a4da1e18 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamRecordBatchList.java +++ b/s3stream/src/main/java/com/automq/stream/utils/biniarysearch/StreamRecordBatchList.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.biniarysearch; diff --git a/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoop.java b/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoop.java index fa7e251e1a..876a8f6128 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoop.java +++ b/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoop.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.threads; diff --git a/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoopSafe.java b/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoopSafe.java new file mode 100644 index 0000000000..8ae8ec2ad4 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/utils/threads/EventLoopSafe.java @@ -0,0 +1,27 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.utils.threads; + +/** + * - If it's annotated to a class means that all the methods in this class should be called from the eventLoop. + * - If it's annotated to a method means that the method should be called in the eventLoop. + */ +public @interface EventLoopSafe { +} diff --git a/s3stream/src/main/java/com/automq/stream/utils/threads/S3StreamThreadPoolMonitor.java b/s3stream/src/main/java/com/automq/stream/utils/threads/S3StreamThreadPoolMonitor.java index 0d83231718..af50abfd78 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/threads/S3StreamThreadPoolMonitor.java +++ b/s3stream/src/main/java/com/automq/stream/utils/threads/S3StreamThreadPoolMonitor.java @@ -1,17 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.threads; import com.automq.stream.utils.ThreadUtils; +import com.automq.stream.utils.Threads; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -20,19 +29,19 @@ import java.util.Collections; import java.util.List; import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; public class S3StreamThreadPoolMonitor { private static final List MONITOR_EXECUTOR = new CopyOnWriteArrayList<>(); - private static final ScheduledExecutorService MONITOR_SCHEDULED = Executors.newSingleThreadScheduledExecutor( - ThreadUtils.createThreadFactory("ThreadPoolMonitor-%d", true)); private static Logger waterMarkLogger = LoggerFactory.getLogger(S3StreamThreadPoolMonitor.class); + private static final ScheduledExecutorService MONITOR_SCHEDULED = + Threads.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("ThreadPoolMonitor-%d", true), waterMarkLogger); private static volatile long threadPoolStatusPeriodTime = TimeUnit.SECONDS.toMillis(3); public static void config(Logger waterMarkLoggerConfig, long threadPoolStatusPeriodTimeConfig) { @@ -89,6 +98,7 @@ public static ThreadPoolExecutor createAndMonitor(int corePoolSize, ThreadUtils.createFastThreadLocalThreadFactory(name + "-%d", isDaemon) : ThreadUtils.createThreadFactory(name + "-%d", isDaemon); + AtomicReference wrapperRef = new AtomicReference<>(); ThreadPoolExecutor executor = new ThreadPoolExecutor( corePoolSize, maximumPoolSize, @@ -102,16 +112,23 @@ protected void afterExecute(Runnable r, Throwable t) { super.afterExecute(r, t); afterExecutionHook.apply(t); } + + @Override + protected void terminated() { + super.terminated(); + ThreadPoolWrapper wrapper = wrapperRef.get(); + if (wrapper != null) { + MONITOR_EXECUTOR.remove(wrapper); + } + } }; List printers = new ArrayList<>(); printers.add(new ThreadPoolQueueSizeMonitor(queueCapacity)); printers.addAll(threadPoolStatusMonitors); - MONITOR_EXECUTOR.add(ThreadPoolWrapper.builder() - .name(name) - .threadPoolExecutor(executor) - .statusPrinters(printers) - .build()); + ThreadPoolWrapper wrapper = ThreadPoolWrapper.builder().name(name).threadPoolExecutor(executor).statusPrinters(printers).build(); + wrapperRef.set(wrapper); + MONITOR_EXECUTOR.add(wrapper); return executor; } diff --git a/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolQueueSizeMonitor.java b/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolQueueSizeMonitor.java index 895ddce75a..fad6378792 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolQueueSizeMonitor.java +++ b/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolQueueSizeMonitor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.threads; diff --git a/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolStatusMonitor.java b/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolStatusMonitor.java index fe62a67e6c..3b6d658c56 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolStatusMonitor.java +++ b/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolStatusMonitor.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.threads; diff --git a/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolWrapper.java b/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolWrapper.java index a21754f149..e7add44d43 100644 --- a/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolWrapper.java +++ b/s3stream/src/main/java/com/automq/stream/utils/threads/ThreadPoolWrapper.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils.threads; diff --git a/s3stream/src/test/java/com/automq/stream/ByteBufSeqAllocTest.java b/s3stream/src/test/java/com/automq/stream/ByteBufSeqAllocTest.java index 50c1c824e8..a061fd643e 100644 --- a/s3stream/src/test/java/com/automq/stream/ByteBufSeqAllocTest.java +++ b/s3stream/src/test/java/com/automq/stream/ByteBufSeqAllocTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/test/java/com/automq/stream/FixedSizeByteBufPoolTest.java b/s3stream/src/test/java/com/automq/stream/FixedSizeByteBufPoolTest.java index b55e2b3a88..c54cf9510d 100644 --- a/s3stream/src/test/java/com/automq/stream/FixedSizeByteBufPoolTest.java +++ b/s3stream/src/test/java/com/automq/stream/FixedSizeByteBufPoolTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream; diff --git a/s3stream/src/test/java/com/automq/stream/s3/AsyncNetworkBandwidthLimiterTest.java b/s3stream/src/test/java/com/automq/stream/s3/AsyncNetworkBandwidthLimiterTest.java index d66585a51f..5045201ca1 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/AsyncNetworkBandwidthLimiterTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/AsyncNetworkBandwidthLimiterTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -85,16 +93,21 @@ public void testThrottleConsume3() { @Test public void testThrottleConsume4() { - AsyncNetworkBandwidthLimiter bucket = new AsyncNetworkBandwidthLimiter(AsyncNetworkBandwidthLimiter.Type.INBOUND, 100, 1000); + AsyncNetworkBandwidthLimiter bucket = new AsyncNetworkBandwidthLimiter( + AsyncNetworkBandwidthLimiter.Type.INBOUND, 100, 1000); bucket.consume(ThrottleStrategy.BYPASS, 1000); Assertions.assertEquals(-100, bucket.getAvailableTokens()); - CompletableFuture cf = bucket.consume(ThrottleStrategy.CATCH_UP, 5); - bucket.consume(ThrottleStrategy.CATCH_UP, 10); - CompletableFuture result = cf.whenComplete((v, e) -> { - Assertions.assertNull(e); - Assertions.assertEquals(95, bucket.getAvailableTokens()); + CompletableFuture firstCompleted = new CompletableFuture<>(); + CompletableFuture cf1 = bucket.consume(ThrottleStrategy.CATCH_UP, 5); + cf1 = cf1.thenApply(v -> { + firstCompleted.complete(true); + return null; + }); + CompletableFuture cf2 = bucket.consume(ThrottleStrategy.CATCH_UP, 10); + CompletableFuture result = cf2.thenAccept(v -> { + Assertions.assertTrue(firstCompleted.isDone(), + "First request should complete before second request"); }); - cf.join(); result.join(); } @@ -105,13 +118,10 @@ public void testThrottleConsume5() throws InterruptedException { Assertions.assertEquals(-200, bucket.getAvailableTokens()); Thread.sleep(500); bucket.consume(ThrottleStrategy.BYPASS, 500); - CompletableFuture cf = bucket.consume(ThrottleStrategy.CATCH_UP, 5); - bucket.consume(ThrottleStrategy.CATCH_UP, 10); - CompletableFuture result = cf.whenComplete((v, e) -> { - Assertions.assertNull(e); - Assertions.assertTrue(bucket.getAvailableTokens() >= 0); - }); - result.join(); + bucket.consume(ThrottleStrategy.CATCH_UP, 5); + CompletableFuture cf = bucket.consume(ThrottleStrategy.CATCH_UP, 10); + cf.join(); + Assertions.assertEquals(-5, bucket.getAvailableTokens()); } @Test diff --git a/s3stream/src/test/java/com/automq/stream/s3/CompositeObjectTest.java b/s3stream/src/test/java/com/automq/stream/s3/CompositeObjectTest.java index 6bf3cde08d..8919d107b5 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/CompositeObjectTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/CompositeObjectTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -122,7 +130,7 @@ public void testCompositeObject_writeAndRead() throws ExecutionException, Interr } StreamRecordBatch newRecord(long streamId, long offset, int count, ByteBuf buf) { - return new StreamRecordBatch(streamId, 0, offset, count, buf); + return StreamRecordBatch.of(streamId, 0, offset, count, buf); } ByteBuf genBuf(byte data, int length) { diff --git a/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatch.java b/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatch.java index 00f2533856..761ed0119f 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatch.java +++ b/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatch.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatchWithContext.java b/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatchWithContext.java index 2a251c6e6e..d7fd213c5c 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatchWithContext.java +++ b/s3stream/src/test/java/com/automq/stream/s3/DefaultRecordBatchWithContext.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/test/java/com/automq/stream/s3/DeltaWALUploadTaskTest.java b/s3stream/src/test/java/com/automq/stream/s3/DefaultUploadWriteAheadLogTaskTest.java similarity index 80% rename from s3stream/src/test/java/com/automq/stream/s3/DeltaWALUploadTaskTest.java rename to s3stream/src/test/java/com/automq/stream/s3/DefaultUploadWriteAheadLogTaskTest.java index e41ba224eb..f538e1f34c 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/DeltaWALUploadTaskTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/DefaultUploadWriteAheadLogTaskTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -49,10 +57,10 @@ import static org.mockito.Mockito.when; @Tag("S3Unit") -public class DeltaWALUploadTaskTest { +public class DefaultUploadWriteAheadLogTaskTest { ObjectManager objectManager; ObjectStorage objectStorage; - DeltaWALUploadTask deltaWALUploadTask; + DefaultUploadWriteAheadLogTask deltaWALUploadTask; @BeforeEach public void setup() { @@ -68,20 +76,20 @@ public void testUpload() throws Exception { Map> map = new HashMap<>(); map.put(233L, List.of( - new StreamRecordBatch(233, 0, 10, 2, random(512)), - new StreamRecordBatch(233, 0, 12, 2, random(128)), - new StreamRecordBatch(233, 0, 14, 2, random(512)) + StreamRecordBatch.of(233, 0, 10, 2, random(512)), + StreamRecordBatch.of(233, 0, 12, 2, random(128)), + StreamRecordBatch.of(233, 0, 14, 2, random(512)) )); map.put(234L, List.of( - new StreamRecordBatch(234, 0, 20, 2, random(128)), - new StreamRecordBatch(234, 0, 22, 2, random(128)) + StreamRecordBatch.of(234, 0, 20, 2, random(128)), + StreamRecordBatch.of(234, 0, 22, 2, random(128)) )); Config config = new Config() .objectBlockSize(16 * 1024 * 1024) .objectPartSize(16 * 1024 * 1024) .streamSplitSize(1000); - deltaWALUploadTask = DeltaWALUploadTask.builder().config(config).streamRecordsMap(map).objectManager(objectManager) + deltaWALUploadTask = DefaultUploadWriteAheadLogTask.builder().config(config).streamRecordsMap(map).objectManager(objectManager) .objectStorage(objectStorage).executor(ForkJoinPool.commonPool()).build(); deltaWALUploadTask.prepare().get(); @@ -153,15 +161,15 @@ public void testUpload_oneStream() throws Exception { Map> map = new HashMap<>(); map.put(233L, List.of( - new StreamRecordBatch(233, 0, 10, 2, random(512)), - new StreamRecordBatch(233, 0, 12, 2, random(128)), - new StreamRecordBatch(233, 0, 14, 2, random(512)) + StreamRecordBatch.of(233, 0, 10, 2, random(512)), + StreamRecordBatch.of(233, 0, 12, 2, random(128)), + StreamRecordBatch.of(233, 0, 14, 2, random(512)) )); Config config = new Config() .objectBlockSize(16 * 1024 * 1024) .objectPartSize(16 * 1024 * 1024) .streamSplitSize(16 * 1024 * 1024); - deltaWALUploadTask = DeltaWALUploadTask.builder().config(config).streamRecordsMap(map).objectManager(objectManager) + deltaWALUploadTask = DefaultUploadWriteAheadLogTask.builder().config(config).streamRecordsMap(map).objectManager(objectManager) .objectStorage(objectStorage).executor(ForkJoinPool.commonPool()).build(); deltaWALUploadTask.prepare().get(); @@ -187,17 +195,17 @@ public void test_emptyWALData() throws ExecutionException, InterruptedException, Map> map = new HashMap<>(); map.put(233L, List.of( - new StreamRecordBatch(233, 0, 10, 2, random(512)) + StreamRecordBatch.of(233, 0, 10, 2, random(512)) )); map.put(234L, List.of( - new StreamRecordBatch(234, 0, 20, 2, random(128)) + StreamRecordBatch.of(234, 0, 20, 2, random(128)) )); Config config = new Config() .objectBlockSize(16 * 1024 * 1024) .objectPartSize(16 * 1024 * 1024) .streamSplitSize(64); - deltaWALUploadTask = DeltaWALUploadTask.builder().config(config).streamRecordsMap(map).objectManager(objectManager) + deltaWALUploadTask = DefaultUploadWriteAheadLogTask.builder().config(config).streamRecordsMap(map).objectManager(objectManager) .objectStorage(objectStorage).executor(ForkJoinPool.commonPool()).build(); assertTrue(deltaWALUploadTask.forceSplit); } diff --git a/s3stream/src/test/java/com/automq/stream/s3/ObjectReaderTest.java b/s3stream/src/test/java/com/automq/stream/s3/ObjectReaderTest.java index 719455b7c5..f83d2136ab 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/ObjectReaderTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/ObjectReaderTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -98,7 +106,7 @@ public void testGetBasicObjectInfo() throws ExecutionException, InterruptedExcep // make index block bigger than 1M int streamCount = 2 * 1024 * 1024 / 40; for (int i = 0; i < streamCount; i++) { - StreamRecordBatch r = new StreamRecordBatch(i, 0, i, 1, TestUtils.random(1)); + StreamRecordBatch r = StreamRecordBatch.of(i, 0, i, 1, TestUtils.random(1)); objectWriter.write(i, List.of(r)); } objectWriter.close().get(); @@ -114,11 +122,11 @@ public void testReadBlockGroup() throws ExecutionException, InterruptedException ObjectStorage objectStorage = new MemoryObjectStorage(); ByteBuf buf = ByteBufAlloc.byteBuffer(0); buf.writeBytes(new ObjectWriter.DataBlock(233L, List.of( - new StreamRecordBatch(233L, 0, 10, 1, TestUtils.random(100)), - new StreamRecordBatch(233L, 0, 11, 2, TestUtils.random(100)) + StreamRecordBatch.of(233L, 0, 10, 1, TestUtils.random(100)), + StreamRecordBatch.of(233L, 0, 11, 2, TestUtils.random(100)) )).buffer()); buf.writeBytes(new ObjectWriter.DataBlock(233L, List.of( - new StreamRecordBatch(233L, 0, 13, 1, TestUtils.random(100)) + StreamRecordBatch.of(233L, 0, 13, 1, TestUtils.random(100)) )).buffer()); int indexPosition = buf.readableBytes(); new DataBlockIndex(233L, 10, 4, 3, 0, buf.readableBytes()).encode(buf); @@ -185,6 +193,6 @@ public void testStreamOffsetRange() throws ExecutionException, InterruptedExcept } StreamRecordBatch newRecord(long streamId, long offset, int count, int payloadSize) { - return new StreamRecordBatch(streamId, 0, offset, count, TestUtils.random(payloadSize)); + return StreamRecordBatch.of(streamId, 0, offset, count, TestUtils.random(payloadSize)); } } diff --git a/s3stream/src/test/java/com/automq/stream/s3/ObjectWriterTest.java b/s3stream/src/test/java/com/automq/stream/s3/ObjectWriterTest.java index d8389b6952..bfd441deff 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/ObjectWriterTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/ObjectWriterTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -129,6 +137,6 @@ public void testWrite_check() { } StreamRecordBatch newRecord(long streamId, long offset, int count, int payloadSize) { - return new StreamRecordBatch(streamId, 0, offset, count, TestUtils.random(payloadSize)); + return StreamRecordBatch.of(streamId, 0, offset, count, TestUtils.random(payloadSize)); } } diff --git a/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java b/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java index aada94055a..58f3c19243 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -16,8 +24,6 @@ import com.automq.stream.s3.cache.blockcache.DefaultObjectReaderFactory; import com.automq.stream.s3.cache.blockcache.StreamReaders; import com.automq.stream.s3.failover.StorageFailureHandler; -import com.automq.stream.s3.metadata.StreamMetadata; -import com.automq.stream.s3.metadata.StreamState; import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.objects.CommitStreamSetObjectRequest; import com.automq.stream.s3.objects.CommitStreamSetObjectResponse; @@ -27,31 +33,38 @@ import com.automq.stream.s3.operator.MemoryObjectStorage; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.streams.StreamManager; +import com.automq.stream.s3.wal.RecordOffset; import com.automq.stream.s3.wal.RecoverResult; import com.automq.stream.s3.wal.WriteAheadLog; import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; import com.automq.stream.s3.wal.impl.MemoryWriteAheadLog; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; -import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; -import io.netty.buffer.ByteBuf; - import static com.automq.stream.s3.TestUtils.random; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; @@ -63,6 +76,8 @@ @Tag("S3Unit") public class S3StorageTest { + private static final Logger LOGGER = LoggerFactory.getLogger(S3StorageTest.class); + StreamManager streamManager; ObjectManager objectManager; WriteAheadLog wal; @@ -71,7 +86,7 @@ public class S3StorageTest { Config config; private static StreamRecordBatch newRecord(long streamId, long offset) { - return new StreamRecordBatch(streamId, 0, offset, 1, random(1)); + return StreamRecordBatch.of(streamId, 0, offset, 1, random(1)); } @BeforeEach @@ -94,13 +109,13 @@ public void testAppend() throws Exception { Mockito.when(objectManager.commitStreamSetObject(any())).thenReturn(CompletableFuture.completedFuture(resp)); CompletableFuture cf1 = storage.append( - new StreamRecordBatch(233, 1, 10, 1, random(100)) + StreamRecordBatch.of(233, 1, 10, 1, random(100)) ); CompletableFuture cf2 = storage.append( - new StreamRecordBatch(233, 1, 11, 2, random(100)) + StreamRecordBatch.of(233, 1, 11, 2, random(100)) ); CompletableFuture cf3 = storage.append( - new StreamRecordBatch(234, 3, 100, 1, random(100)) + StreamRecordBatch.of(234, 3, 100, 1, random(100)) ); cf1.get(3, TimeUnit.SECONDS); @@ -127,58 +142,6 @@ public void testAppend() throws Exception { assertEquals(101, streamRanges.get(1).getEndOffset()); } - @Test - public void testWALConfirmOffsetCalculator() { - S3Storage.WALConfirmOffsetCalculator calc = new S3Storage.WALConfirmOffsetCalculator(); - WalWriteRequest r0 = new WalWriteRequest(null, 0L, null); - WalWriteRequest r1 = new WalWriteRequest(null, 1L, null); - WalWriteRequest r2 = new WalWriteRequest(null, 2L, null); - WalWriteRequest r3 = new WalWriteRequest(null, 3L, null); - - calc.add(r3); - calc.add(r1); - calc.add(r2); - calc.add(r0); - - calc.update(); - assertEquals(-1L, calc.get()); - - r0.confirmed = true; - calc.update(); - assertEquals(0L, calc.get()); - - r3.confirmed = true; - calc.update(); - assertEquals(0L, calc.get()); - - r1.confirmed = true; - calc.update(); - assertEquals(1L, calc.get()); - - r2.confirmed = true; - calc.update(); - assertEquals(3L, calc.get()); - } - - @Test - public void testWALCallbackSequencer() { - S3Storage.WALCallbackSequencer seq = new S3Storage.WALCallbackSequencer(); - WalWriteRequest r0 = new WalWriteRequest(newRecord(233L, 10L), 100L, new CompletableFuture<>()); - WalWriteRequest r1 = new WalWriteRequest(newRecord(233L, 11L), 101L, new CompletableFuture<>()); - WalWriteRequest r2 = new WalWriteRequest(newRecord(234L, 20L), 102L, new CompletableFuture<>()); - WalWriteRequest r3 = new WalWriteRequest(newRecord(234L, 21L), 103L, new CompletableFuture<>()); - - seq.before(r0); - seq.before(r1); - seq.before(r2); - seq.before(r3); - - assertEquals(Collections.emptyList(), seq.after(r3)); - assertEquals(List.of(r2, r3), seq.after(r2)); - assertEquals(List.of(r0), seq.after(r0)); - assertEquals(List.of(r1), seq.after(r1)); - } - @Test public void testUploadWALObject_sequence() throws ExecutionException, InterruptedException, TimeoutException { List> objectIdCfList = List.of(new CompletableFuture<>(), new CompletableFuture<>()); @@ -192,13 +155,13 @@ public void testUploadWALObject_sequence() throws ExecutionException, Interrupte LogCache.LogCacheBlock logCacheBlock1 = new LogCache.LogCacheBlock(1024); logCacheBlock1.put(newRecord(233L, 10L)); logCacheBlock1.put(newRecord(234L, 10L)); - logCacheBlock1.confirmOffset(10L); + logCacheBlock1.lastRecordOffset(DefaultRecordOffset.of(0, 10L, 0)); CompletableFuture cf1 = storage.uploadDeltaWAL(logCacheBlock1); LogCache.LogCacheBlock logCacheBlock2 = new LogCache.LogCacheBlock(1024); logCacheBlock2.put(newRecord(233L, 20L)); logCacheBlock2.put(newRecord(234L, 20L)); - logCacheBlock2.confirmOffset(20L); + logCacheBlock2.lastRecordOffset(DefaultRecordOffset.of(0, 20L, 0)); CompletableFuture cf2 = storage.uploadDeltaWAL(logCacheBlock2); // sequence get objectId @@ -223,48 +186,59 @@ public void testUploadWALObject_sequence() throws ExecutionException, Interrupte @Test public void testRecoverContinuousRecords() { List recoverResults = List.of( - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(233L, 10L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(233L, 11L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(233L, 12L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(233L, 15L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(234L, 20L))) + new TestRecoverResult(newRecord(233L, 10L)), + new TestRecoverResult(newRecord(233L, 11L)), + new TestRecoverResult(newRecord(233L, 12L)), + new TestRecoverResult(newRecord(233L, 15L)), + new TestRecoverResult(newRecord(234L, 20L)) ); + Iterator iterator = recoverResults.iterator(); - List openingStreams = List.of(new StreamMetadata(233L, 0L, 0L, 11L, StreamState.OPENED)); - LogCache.LogCacheBlock cacheBlock = S3Storage.recoverContinuousRecords(recoverResults.iterator(), openingStreams); + Map streamEndOffsets = Map.of(233L, 11L); + S3Storage.RecoveryBlockResult result = S3Storage.recoverContinuousRecords(iterator, streamEndOffsets, 1 << 30, LOGGER); + LogCache.LogCacheBlock cacheBlock = result.cacheBlock; + assertNull(result.exception); // ignore closed stream and noncontinuous records. assertEquals(1, cacheBlock.records().size()); List streamRecords = cacheBlock.records().get(233L); assertEquals(2, streamRecords.size()); assertEquals(11L, streamRecords.get(0).getBaseOffset()); assertEquals(12L, streamRecords.get(1).getBaseOffset()); + } + + @Test + public void testRecoverDataLoss() { + List recoverResults = List.of( + new TestRecoverResult(newRecord(233L, 10L)), + new TestRecoverResult(newRecord(233L, 11L)), + new TestRecoverResult(newRecord(233L, 12L)) + ); + Iterator iterator = recoverResults.iterator(); // simulate data loss - openingStreams = List.of( - new StreamMetadata(233L, 0L, 0L, 5L, StreamState.OPENED)); - boolean exception = false; - try { - S3Storage.recoverContinuousRecords(recoverResults.iterator(), openingStreams); - } catch (IllegalStateException e) { - exception = true; - } - Assertions.assertTrue(exception); + Map streamEndOffsets = Map.of(233L, 5L); + S3Storage.RecoveryBlockResult result = S3Storage.recoverContinuousRecords(iterator, streamEndOffsets, 1 << 30, LOGGER); + assertNotNull(result.exception); + LogCache.LogCacheBlock cacheBlock = result.cacheBlock; + assertEquals(0, cacheBlock.records().size()); } @Test public void testRecoverOutOfOrderRecords() { List recoverResults = List.of( - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 9L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 10L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 13L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 11L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 12L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 14L))), - new TestRecoverResult(StreamRecordBatchCodec.encode(newRecord(42L, 20L))) + new TestRecoverResult(newRecord(42L, 9L)), + new TestRecoverResult(newRecord(42L, 10L)), + new TestRecoverResult(newRecord(42L, 13L)), + new TestRecoverResult(newRecord(42L, 11L)), + new TestRecoverResult(newRecord(42L, 12L)), + new TestRecoverResult(newRecord(42L, 14L)), + new TestRecoverResult(newRecord(42L, 20L)) ); - List openingStreams = List.of(new StreamMetadata(42L, 0L, 0L, 10L, StreamState.OPENED)); - LogCache.LogCacheBlock cacheBlock = S3Storage.recoverContinuousRecords(recoverResults.iterator(), openingStreams); + Map streamEndOffsets = Map.of(42L, 10L); + S3Storage.RecoveryBlockResult result = S3Storage.recoverContinuousRecords(recoverResults.iterator(), streamEndOffsets, 1 << 30, LOGGER); + LogCache.LogCacheBlock cacheBlock = result.cacheBlock; + assertNull(result.exception); // ignore closed stream and noncontinuous records. assertEquals(1, cacheBlock.records().size()); List streamRecords = cacheBlock.records().get(42L); @@ -276,6 +250,50 @@ public void testRecoverOutOfOrderRecords() { assertEquals(14L, streamRecords.get(4).getBaseOffset()); } + @Test + public void testSegmentedRecovery() { + List recoverResults = List.of( + new TestRecoverResult(newRecord(42L, 10L)), + new TestRecoverResult(newRecord(42L, 11L)), + new TestRecoverResult(newRecord(42L, 12L)), + new TestRecoverResult(newRecord(42L, 13L)), + new TestRecoverResult(newRecord(42L, 14L)), + new TestRecoverResult(newRecord(42L, 20L)) + ); + Iterator iterator = recoverResults.iterator(); + + Map streamEndOffsets = new HashMap<>(); + S3Storage.RecoveryBlockResult result; + List streamRecords; + final long maxCacheSize = 200L; + + streamEndOffsets.put(42L, 10L); + result = S3Storage.recoverContinuousRecords(iterator, streamEndOffsets, maxCacheSize, LOGGER); + assertNull(result.exception); + assertTrue(result.cacheBlock.isFull()); + streamRecords = result.cacheBlock.records().get(42L); + assertEquals(2, streamRecords.size()); + assertEquals(10L, streamRecords.get(0).getBaseOffset()); + assertEquals(11L, streamRecords.get(1).getBaseOffset()); + + streamEndOffsets.put(42L, 12L); + result = S3Storage.recoverContinuousRecords(iterator, streamEndOffsets, maxCacheSize, LOGGER); + assertNull(result.exception); + assertTrue(result.cacheBlock.isFull()); + streamRecords = result.cacheBlock.records().get(42L); + assertEquals(2, streamRecords.size()); + assertEquals(12L, streamRecords.get(0).getBaseOffset()); + assertEquals(13L, streamRecords.get(1).getBaseOffset()); + + streamEndOffsets.put(42L, 14L); + result = S3Storage.recoverContinuousRecords(iterator, streamEndOffsets, maxCacheSize, LOGGER); + assertNull(result.exception); + assertFalse(result.cacheBlock.isFull()); + streamRecords = result.cacheBlock.records().get(42L); + assertEquals(1, streamRecords.size()); + assertEquals(14L, streamRecords.get(0).getBaseOffset()); + } + @Test public void testWALOverCapacity() throws OverCapacityException { storage.append(newRecord(233L, 10L)); @@ -300,20 +318,20 @@ public void testWALOverCapacity() throws OverCapacityException { } static class TestRecoverResult implements RecoverResult { - private final ByteBuf record; + private final StreamRecordBatch record; - public TestRecoverResult(ByteBuf record) { + public TestRecoverResult(StreamRecordBatch record) { this.record = record; } @Override - public ByteBuf record() { + public StreamRecordBatch record() { return record; } @Override - public long recordOffset() { - return 0; + public RecordOffset recordOffset() { + return DefaultRecordOffset.of(0, 0, 0); } } -} +} \ No newline at end of file diff --git a/s3stream/src/test/java/com/automq/stream/s3/S3StreamClientTest.java b/s3stream/src/test/java/com/automq/stream/s3/S3StreamClientTest.java index 89e40a878f..897af81a4b 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/S3StreamClientTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/S3StreamClientTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -67,7 +75,7 @@ public void testShutdown_withOpeningStream() { CompletableFuture cf = new CompletableFuture<>(); when(streamManager.openStream(anyLong(), anyLong(), anyMap())).thenReturn(cf); - doAnswer(args -> stream).when(client).newStream(any()); + doAnswer(args -> stream).when(client).newStream(any(), any()); scheduler.schedule(() -> { cf.complete(new StreamMetadata(1, 2, 100, 200, StreamState.OPENED)); diff --git a/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java b/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java index 8e94da73a9..9e1cfdc01c 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; @@ -44,7 +52,7 @@ public class S3StreamTest { public void setup() { storage = mock(Storage.class); streamManager = mock(StreamManager.class); - stream = new S3Stream(233, 1, 100, 233, storage, streamManager); + stream = S3Stream.create(233L, 1L, 100L, 233L, storage, streamManager); } @Test @@ -72,7 +80,7 @@ public void testFetch() throws Throwable { } ReadDataBlock newReadDataBlock(long start, long end, int size) { - StreamRecordBatch record = new StreamRecordBatch(0, 0, start, (int) (end - start), TestUtils.random(size)); + StreamRecordBatch record = StreamRecordBatch.of(0, 0, start, (int) (end - start), TestUtils.random(size)); return new ReadDataBlock(List.of(record), CacheAccessType.DELTA_WAL_CACHE_HIT); } } diff --git a/s3stream/src/test/java/com/automq/stream/s3/TestUtils.java b/s3stream/src/test/java/com/automq/stream/s3/TestUtils.java index 144c0e722d..f011cfface 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/TestUtils.java +++ b/s3stream/src/test/java/com/automq/stream/s3/TestUtils.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3; diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/AsyncLRUCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/AsyncLRUCacheTest.java index 7739d182fc..762909dd58 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/AsyncLRUCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/AsyncLRUCacheTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/LRUCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/LRUCacheTest.java index f8ff15b437..22f96edfe2 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/LRUCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/LRUCacheTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/LogCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/LogCacheTest.java index d64644d28f..e1fe8c25e4 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/LogCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/LogCacheTest.java @@ -1,17 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; import com.automq.stream.s3.TestUtils; +import com.automq.stream.s3.cache.LogCache.LogCacheBlock; import com.automq.stream.s3.model.StreamRecordBatch; import org.junit.jupiter.api.Tag; @@ -19,8 +28,11 @@ import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutionException; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; @Tag("S3Unit") public class LogCacheTest { @@ -29,15 +41,15 @@ public class LogCacheTest { public void testPutGet() { LogCache logCache = new LogCache(1024 * 1024, 1024 * 1024); - logCache.put(new StreamRecordBatch(233L, 0L, 10L, 1, TestUtils.random(20))); - logCache.put(new StreamRecordBatch(233L, 0L, 11L, 2, TestUtils.random(20))); + logCache.put(StreamRecordBatch.of(233L, 0L, 10L, 1, TestUtils.random(20))); + logCache.put(StreamRecordBatch.of(233L, 0L, 11L, 2, TestUtils.random(20))); logCache.archiveCurrentBlock(); - logCache.put(new StreamRecordBatch(233L, 0L, 13L, 2, TestUtils.random(20))); + logCache.put(StreamRecordBatch.of(233L, 0L, 13L, 2, TestUtils.random(20))); logCache.archiveCurrentBlock(); - logCache.put(new StreamRecordBatch(233L, 0L, 20L, 1, TestUtils.random(20))); - logCache.put(new StreamRecordBatch(233L, 0L, 21L, 1, TestUtils.random(20))); + logCache.put(StreamRecordBatch.of(233L, 0L, 20L, 1, TestUtils.random(20))); + logCache.put(StreamRecordBatch.of(233L, 0L, 21L, 1, TestUtils.random(20))); List records = logCache.get(233L, 10L, 21L, 1000); assertEquals(1, records.size()); @@ -62,7 +74,7 @@ public void testOffsetIndex() { LogCache cache = new LogCache(Integer.MAX_VALUE, Integer.MAX_VALUE); for (int i = 0; i < 100000; i++) { - cache.put(new StreamRecordBatch(233L, 0L, i, 1, TestUtils.random(1))); + cache.put(StreamRecordBatch.of(233L, 0L, i, 1, TestUtils.random(1))); } long start = System.nanoTime(); @@ -75,4 +87,122 @@ public void testOffsetIndex() { assertEquals(100000, offsetIndexMap.get(100000L).index); } + @Test + public void testClearStreamRecords() { + LogCache logCache = new LogCache(1024 * 1024, 1024 * 1024); + + logCache.put(StreamRecordBatch.of(233L, 0L, 10L, 1, TestUtils.random(20))); + logCache.put(StreamRecordBatch.of(233L, 0L, 11L, 2, TestUtils.random(20))); + + logCache.archiveCurrentBlock(); + logCache.put(StreamRecordBatch.of(233L, 0L, 13L, 2, TestUtils.random(20))); + + logCache.put(StreamRecordBatch.of(234L, 0L, 13L, 2, TestUtils.random(20))); + + assertTrue(logCache.blocks.get(0).containsStream(233L)); + assertTrue(logCache.blocks.get(1).containsStream(234L)); + logCache.clearStreamRecords(233L); + assertFalse(logCache.blocks.get(0).containsStream(233L)); + assertTrue(logCache.blocks.get(1).containsStream(234L)); + + logCache.clearStreamRecords(234L); + assertFalse(logCache.blocks.get(0).containsStream(233L)); + assertFalse(logCache.blocks.get(1).containsStream(234L)); + } + + @Test + public void testIsDiscontinuous() { + LogCacheBlock left = new LogCacheBlock(1024L * 1024); + left.put(StreamRecordBatch.of(233L, 0L, 10L, 1, TestUtils.random(20))); + + LogCacheBlock right = new LogCacheBlock(1024L * 1024); + right.put(StreamRecordBatch.of(233L, 0L, 13L, 1, TestUtils.random(20))); + + assertTrue(LogCache.isDiscontinuous(left, right)); + + left = new LogCacheBlock(1024L * 1024); + left.put(StreamRecordBatch.of(233L, 0L, 10L, 1, TestUtils.random(20))); + left.put(StreamRecordBatch.of(234L, 0L, 10L, 1, TestUtils.random(20))); + + right = new LogCacheBlock(1024L * 1024); + right.put(StreamRecordBatch.of(233L, 0L, 11L, 1, TestUtils.random(20))); + assertFalse(LogCache.isDiscontinuous(left, right)); + } + + @Test + public void testMergeBlock() { + long size = 0; + LogCacheBlock left = new LogCacheBlock(1024L * 1024); + left.put(StreamRecordBatch.of(233L, 0L, 10L, 1, TestUtils.random(20))); + left.put(StreamRecordBatch.of(234L, 0L, 100L, 1, TestUtils.random(20))); + size += left.size(); + + LogCacheBlock right = new LogCacheBlock(1024L * 1024); + right.put(StreamRecordBatch.of(233L, 0L, 11L, 1, TestUtils.random(20))); + right.put(StreamRecordBatch.of(235L, 0L, 200L, 1, TestUtils.random(20))); + size += right.size(); + + LogCache.mergeBlock(left, right); + assertEquals(size, left.size()); + LogCache.StreamCache stream233 = left.map.get(233L); + assertEquals(10, stream233.startOffset()); + assertEquals(12, stream233.endOffset()); + assertEquals(2, stream233.records.size()); + assertEquals(10, stream233.records.get(0).getBaseOffset()); + assertEquals(11, stream233.records.get(1).getBaseOffset()); + + LogCache.StreamCache stream234 = left.map.get(234L); + assertEquals(100, stream234.startOffset()); + assertEquals(101, stream234.endOffset()); + assertEquals(1, stream234.records.size()); + assertEquals(100, stream234.records.get(0).getBaseOffset()); + + LogCache.StreamCache stream235 = left.map.get(235L); + assertEquals(200, stream235.startOffset()); + assertEquals(201, stream235.endOffset()); + assertEquals(1, stream235.records.size()); + assertEquals(200, stream235.records.get(0).getBaseOffset()); + } + + @Test + public void testTryMergeLogic() throws ExecutionException, InterruptedException { + LogCache logCache = new LogCache(Long.MAX_VALUE, 10_000L); + final long streamId = 233L; + final int blocksToCreate = LogCache.MERGE_BLOCK_THRESHOLD + 2; + + // create multiple blocks, each containing one record for the same stream with contiguous offsets + for (int i = 0; i < blocksToCreate; i++) { + logCache.put(StreamRecordBatch.of(streamId, 0L, i, 1, TestUtils.random(1))); + logCache.archiveCurrentBlock(); + } + + int before = logCache.blocks.size(); + assertTrue(before > LogCache.MERGE_BLOCK_THRESHOLD, "need more than 8 blocks to exercise tryMerge"); + + LogCache.LogCacheBlock left = logCache.blocks.get(0); + LogCache.LogCacheBlock right = logCache.blocks.get(1); + + // verify contiguous condition before merge: left.end == right.start + LogCache.StreamCache leftCache = left.map.get(streamId); + LogCache.StreamCache rightCache = right.map.get(streamId); + assertEquals(leftCache.endOffset(), rightCache.startOffset()); + + // mark both blocks free to trigger tryMerge (called inside markFree) + logCache.markFree(left).get(); + logCache.markFree(right).get(); + + int after = logCache.blocks.size(); + assertEquals(before - 1, after, "two adjacent free contiguous blocks should be merged into one"); + + // verify merged block contains both records and correct range + LogCache.LogCacheBlock merged = logCache.blocks.get(0); + assertTrue(merged.free); + LogCache.StreamCache mergedCache = merged.map.get(streamId); + assertEquals(2, mergedCache.records.size()); + assertEquals(0L, mergedCache.startOffset()); + assertEquals(2L, mergedCache.endOffset()); + assertEquals(0L, mergedCache.records.get(0).getBaseOffset()); + assertEquals(1L, mergedCache.records.get(1).getBaseOffset()); + } + } diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/ObjectReaderLRUCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/ObjectReaderLRUCacheTest.java index ea207bbb16..8cd84bd75f 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/ObjectReaderLRUCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/ObjectReaderLRUCacheTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache; @@ -38,7 +46,7 @@ public class ObjectReaderLRUCacheTest { private void writeStream(int streamCount, ObjectWriter objectWriter) { for (int i = 0; i < streamCount; i++) { - StreamRecordBatch r = new StreamRecordBatch(i, 0, i, 1, TestUtils.random(1)); + StreamRecordBatch r = StreamRecordBatch.of(i, 0, i, 1, TestUtils.random(1)); objectWriter.write(i, List.of(r)); } } diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/DataBlockCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/DataBlockCacheTest.java index b8b256b498..dcd739f55a 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/DataBlockCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/DataBlockCacheTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -217,10 +225,10 @@ private ByteBuf newDataBlockGroupBuf(DataBlockIndex index) { long offset = index.startOffset(); // the first N - 1 record's count = 1, body size = 1 for (int i = 0; i < index.recordCount() - 1; i++, offset++) { - records.add(new StreamRecordBatch(STREAM_ID, 0, offset, 1, TestUtils.random(1))); + records.add(StreamRecordBatch.of(STREAM_ID, 0, offset, 1, TestUtils.random(1))); } // the last record padding the remaining - records.add(new StreamRecordBatch(STREAM_ID, 0, offset, index.endOffsetDelta() - (index.recordCount() - 1), TestUtils.random(remainingSize))); + records.add(StreamRecordBatch.of(STREAM_ID, 0, offset, index.endOffsetDelta() - (index.recordCount() - 1), TestUtils.random(remainingSize))); ByteBuf buf = new ObjectWriter.DataBlock(STREAM_ID, records).buffer(); assertEquals(index.size(), buf.readableBytes()); return buf; diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/MockObject.java b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/MockObject.java index 884106fb78..8ef0394c88 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/MockObject.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/MockObject.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -18,6 +26,7 @@ import com.automq.stream.s3.metadata.StreamOffsetRange; import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.operator.MemoryObjectStorage; +import com.automq.stream.s3.operator.ObjectStorage; import java.util.LinkedList; import java.util.List; @@ -53,7 +62,7 @@ public static class Builder { public Builder(long objectId, int blockSizeThreshold) { this.objectId = objectId; - this.writer = new ObjectWriter.DefaultObjectWriter(objectId, operator, blockSizeThreshold, Integer.MAX_VALUE); + this.writer = new ObjectWriter.DefaultObjectWriter(objectId, operator, blockSizeThreshold, Integer.MAX_VALUE, new ObjectStorage.WriteOptions()); } public Builder mockDelay(long delay) { diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReaderTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReaderTest.java index e98189da5a..a995863bc1 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReaderTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReaderTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.cache.blockcache; @@ -82,21 +90,21 @@ void setup() { // object=6 [24, 29) // object=7 [29, 34) objects.put(0L, MockObject.builder(0, BLOCK_SIZE_THRESHOLD).mockDelay(100).write(STREAM_ID, List.of( - new StreamRecordBatch(STREAM_ID, 0, 0, 1, TestUtils.random(1)) + StreamRecordBatch.of(STREAM_ID, 0, 0, 1, TestUtils.random(1)) )).build()); objects.put(1L, MockObject.builder(1L, 1).mockDelay(100).write(STREAM_ID, List.of( - new StreamRecordBatch(STREAM_ID, 0, 1, 1, TestUtils.random(19)), - new StreamRecordBatch(STREAM_ID, 0, 2, 1, TestUtils.random(10)), - new StreamRecordBatch(STREAM_ID, 0, 3, 1, TestUtils.random(10)) + StreamRecordBatch.of(STREAM_ID, 0, 1, 1, TestUtils.random(19)), + StreamRecordBatch.of(STREAM_ID, 0, 2, 1, TestUtils.random(10)), + StreamRecordBatch.of(STREAM_ID, 0, 3, 1, TestUtils.random(10)) )).build()); for (int i = 0; i < 6; i++) { long offset = 4 + i * 5; objects.put(i + 2L, MockObject.builder(i + 2L, BLOCK_SIZE_THRESHOLD).mockDelay(100).write(STREAM_ID, List.of( - new StreamRecordBatch(STREAM_ID, 0, offset, 1, TestUtils.random(1024 * 1024 / 4)), - new StreamRecordBatch(STREAM_ID, 0, offset + 1, 1, TestUtils.random(1024 * 1024 / 4)), - new StreamRecordBatch(STREAM_ID, 0, offset + 2, 1, TestUtils.random(1024 * 1024 / 4)), - new StreamRecordBatch(STREAM_ID, 0, offset + 3, 1, TestUtils.random(1024 * 1024 / 4)), - new StreamRecordBatch(STREAM_ID, 0, offset + 4, 1, TestUtils.random(1024 * 1024 / 4)) + StreamRecordBatch.of(STREAM_ID, 0, offset, 1, TestUtils.random(1024 * 1024 / 4)), + StreamRecordBatch.of(STREAM_ID, 0, offset + 1, 1, TestUtils.random(1024 * 1024 / 4)), + StreamRecordBatch.of(STREAM_ID, 0, offset + 2, 1, TestUtils.random(1024 * 1024 / 4)), + StreamRecordBatch.of(STREAM_ID, 0, offset + 3, 1, TestUtils.random(1024 * 1024 / 4)), + StreamRecordBatch.of(STREAM_ID, 0, offset + 4, 1, TestUtils.random(1024 * 1024 / 4)) )).build()); } diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReadersTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReadersTest.java new file mode 100644 index 0000000000..3b81ddd5b4 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/cache/blockcache/StreamReadersTest.java @@ -0,0 +1,190 @@ +/* + * Copyright 2024, AutoMQ HK Limited. + * + * The use of this file is governed by the Business Source License, + * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * + * As of the Change Date specified in that file, in accordance with + * the Business Source License, use of this software will be governed + * by the Apache License, Version 2.0 + */ + +package com.automq.stream.s3.cache.blockcache; + +import com.automq.stream.s3.ObjectReader; +import com.automq.stream.s3.TestUtils; +import com.automq.stream.s3.cache.ReadDataBlock; +import com.automq.stream.s3.metadata.S3ObjectMetadata; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.objects.ObjectManager; +import com.automq.stream.s3.operator.MemoryObjectStorage; +import com.automq.stream.s3.operator.ObjectStorage; +import com.automq.stream.s3.trace.context.TraceContext; +import com.automq.stream.utils.MockTime; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +import static org.awaitility.Awaitility.await; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +public class StreamReadersTest { + private static final long STREAM_ID_1 = 100L; + private static final long STREAM_ID_2 = 200L; + private static final int BLOCK_SIZE_THRESHOLD = 1024; + + private Map objects; + private ObjectManager objectManager; + private ObjectStorage objectStorage; + private ObjectReaderFactory objectReaderFactory; + private StreamReaders streamReaders; + private MockTime mockTime; + + @BeforeEach + void setup() { + objects = new HashMap<>(); + + // Create mock objects for testing with different offset ranges + // Object 1: STREAM_ID_1 offset 0-2 + objects.put(1L, MockObject.builder(1L, BLOCK_SIZE_THRESHOLD).write(STREAM_ID_1, List.of( + StreamRecordBatch.of(STREAM_ID_1, 0, 0, 2, TestUtils.random(100)) + )).build()); + // Object 2: STREAM_ID_2 offset 0-1 + objects.put(2L, MockObject.builder(2L, BLOCK_SIZE_THRESHOLD).write(STREAM_ID_2, List.of( + StreamRecordBatch.of(STREAM_ID_2, 0, 0, 1, TestUtils.random(100)) + )).build()); + + objectManager = mock(ObjectManager.class); + + when(objectManager.isObjectExist(anyLong())).thenReturn(true); + // Mock getObjects method to return appropriate objects based on offset ranges + // For STREAM_ID_1, use the combined object that covers 0-2 range + when(objectManager.getObjects(eq(STREAM_ID_1), anyLong(), anyLong(), anyInt())) + .thenReturn(CompletableFuture.completedFuture(List.of(objects.get(1L).metadata))); + // STREAM_ID_2 offset 0-1 -> object 3 + when(objectManager.getObjects(eq(STREAM_ID_2), anyLong(), anyLong(), anyInt())) + .thenReturn(CompletableFuture.completedFuture(List.of(objects.get(2L).metadata))); + + objectStorage = new MemoryObjectStorage(); + + objectReaderFactory = new ObjectReaderFactory() { + @Override + public ObjectReader get(S3ObjectMetadata metadata) { + return objects.get(metadata.objectId()).objectReader(); + } + + @Override + public ObjectStorage getObjectStorage() { + return objectStorage; + } + }; + + mockTime = new MockTime(); + streamReaders = new StreamReaders(Long.MAX_VALUE, objectManager, objectStorage, objectReaderFactory, 2, mockTime); + } + + @AfterEach + void tearDown() { + if (streamReaders != null) { + // Clean up resources + streamReaders = null; + } + } + + @Test + public void testStreamReaderCreationAndReuse() throws Exception { + TraceContext context = TraceContext.DEFAULT; + + // Initially no StreamReaders + assertEquals(0, streamReaders.getActiveStreamReaderCount()); + + // Create first StreamReader + CompletableFuture readFuture1 = streamReaders.read(context, STREAM_ID_1, 0, 1, Integer.MAX_VALUE); + ReadDataBlock result1 = readFuture1.get(5, TimeUnit.SECONDS); + result1.getRecords().forEach(StreamRecordBatch::release); + + assertEquals(1, streamReaders.getActiveStreamReaderCount()); + + // Read from same stream again - should reuse existing StreamReader + CompletableFuture readFuture2 = streamReaders.read(context, STREAM_ID_1, 1, 2, Integer.MAX_VALUE); + ReadDataBlock result2 = readFuture2.get(5, TimeUnit.SECONDS); + result2.getRecords().forEach(StreamRecordBatch::release); + + // Should still have 1 StreamReader (reused) + assertEquals(1, streamReaders.getActiveStreamReaderCount()); + } + + @Test + public void testCleanupTrigger() throws Exception { + TraceContext context = TraceContext.DEFAULT; + + // Create some StreamReaders + CompletableFuture readFuture1 = streamReaders.read(context, STREAM_ID_1, 0, 1, Integer.MAX_VALUE); + ReadDataBlock result1 = readFuture1.get(5, TimeUnit.SECONDS); + result1.getRecords().forEach(StreamRecordBatch::release); + + CompletableFuture readFuture2 = streamReaders.read(context, STREAM_ID_2, 0, 1, Integer.MAX_VALUE); + ReadDataBlock result2 = readFuture2.get(5, TimeUnit.SECONDS); + result2.getRecords().forEach(StreamRecordBatch::release); + + assertEquals(2, streamReaders.getActiveStreamReaderCount()); + + // Trigger cleanup - should not affect non-expired readers + streamReaders.triggerExpiredStreamReaderCleanup(); + + // Wait for async cleanup to complete + await().atMost(1, TimeUnit.SECONDS) + .pollInterval(100, TimeUnit.MILLISECONDS) + .until(() -> streamReaders.getActiveStreamReaderCount() == 2); + + // StreamReaders should still be there (not expired yet) + assertEquals(2, streamReaders.getActiveStreamReaderCount()); + } + + @Test + public void testExpiredStreamReaderCleanupExecution() throws Exception { + TraceContext context = TraceContext.DEFAULT; + + // Create a StreamReader + CompletableFuture readFuture = streamReaders.read(context, STREAM_ID_1, 0, 1, Integer.MAX_VALUE); + ReadDataBlock result = readFuture.get(5, TimeUnit.SECONDS); + result.getRecords().forEach(StreamRecordBatch::release); + + assertEquals(1, streamReaders.getActiveStreamReaderCount()); + + // Advance mock time to simulate expiration (advance by 2 minutes, expiration is 1 minute) + mockTime.sleep(TimeUnit.MINUTES.toMillis(2)); + + // Trigger cleanup - should now clean up expired StreamReaders + streamReaders.triggerExpiredStreamReaderCleanup(); + + // Wait for async cleanup to complete + await().atMost(5, TimeUnit.SECONDS) + .pollInterval(100, TimeUnit.MILLISECONDS) + .until(() -> streamReaders.getActiveStreamReaderCount() == 0); + + // Verify system still works after cleanup + CompletableFuture readFuture2 = streamReaders.read(context, STREAM_ID_2, 0, 1, Integer.MAX_VALUE); + ReadDataBlock result2 = readFuture2.get(5, TimeUnit.SECONDS); + result2.getRecords().forEach(StreamRecordBatch::release); + + assertEquals(1, streamReaders.getActiveStreamReaderCount()); + } + + + +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionAnalyzerTest.java b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionAnalyzerTest.java index f6fcf1b4e7..2698d904fe 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionAnalyzerTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionAnalyzerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionManagerTest.java b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionManagerTest.java index c6f586379d..4ab3ee5846 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionManagerTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; @@ -60,8 +68,13 @@ import io.netty.buffer.ByteBuf; import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.http.HttpStatusCode; import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest; +import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse; import software.amazon.awssdk.services.s3.model.PutObjectRequest; +import software.amazon.awssdk.services.s3.model.S3Exception; +import software.amazon.awssdk.services.s3.model.UploadPartRequest; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -337,9 +350,9 @@ public void testCompactSingleObject() { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_3, objectId); ObjectWriter objectWriter = ObjectWriter.writer(OBJECT_3, objectStorage, 1024, 1024); - StreamRecordBatch r1 = new StreamRecordBatch(STREAM_1, 0, 500, 20, TestUtils.random(20)); - StreamRecordBatch r2 = new StreamRecordBatch(STREAM_3, 0, 0, 10, TestUtils.random(1024)); - StreamRecordBatch r3 = new StreamRecordBatch(STREAM_3, 0, 10, 10, TestUtils.random(1024)); + StreamRecordBatch r1 = StreamRecordBatch.of(STREAM_1, 0, 500, 20, TestUtils.random(20)); + StreamRecordBatch r2 = StreamRecordBatch.of(STREAM_3, 0, 0, 10, TestUtils.random(1024)); + StreamRecordBatch r3 = StreamRecordBatch.of(STREAM_3, 0, 10, 10, TestUtils.random(1024)); objectWriter.write(STREAM_1, List.of(r1)); objectWriter.write(STREAM_3, List.of(r2, r3)); objectWriter.close().join(); @@ -417,9 +430,9 @@ public void testCompactionWithDataTrimmed3() { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_3, objectId); ObjectWriter objectWriter = ObjectWriter.writer(OBJECT_3, objectStorage, 1024, 1024); - StreamRecordBatch r1 = new StreamRecordBatch(STREAM_1, 0, 500, 20, TestUtils.random(20)); - StreamRecordBatch r2 = new StreamRecordBatch(STREAM_3, 0, 0, 10, TestUtils.random(1024)); - StreamRecordBatch r3 = new StreamRecordBatch(STREAM_3, 0, 10, 10, TestUtils.random(1024)); + StreamRecordBatch r1 = StreamRecordBatch.of(STREAM_1, 0, 500, 20, TestUtils.random(20)); + StreamRecordBatch r2 = StreamRecordBatch.of(STREAM_3, 0, 0, 10, TestUtils.random(1024)); + StreamRecordBatch r3 = StreamRecordBatch.of(STREAM_3, 0, 10, 10, TestUtils.random(1024)); objectWriter.write(STREAM_1, List.of(r1)); objectWriter.write(STREAM_3, List.of(r2, r3)); objectWriter.close().join(); @@ -449,9 +462,9 @@ public void testCompactionWithDataTrimmed4() { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_3, objectId); ObjectWriter objectWriter = ObjectWriter.writer(OBJECT_3, objectStorage, 200, 1024); - StreamRecordBatch r1 = new StreamRecordBatch(STREAM_1, 0, 500, 20, TestUtils.random(20)); - StreamRecordBatch r2 = new StreamRecordBatch(STREAM_3, 0, 0, 10, TestUtils.random(200)); - StreamRecordBatch r3 = new StreamRecordBatch(STREAM_3, 0, 10, 10, TestUtils.random(200)); + StreamRecordBatch r1 = StreamRecordBatch.of(STREAM_1, 0, 500, 20, TestUtils.random(20)); + StreamRecordBatch r2 = StreamRecordBatch.of(STREAM_3, 0, 0, 10, TestUtils.random(200)); + StreamRecordBatch r3 = StreamRecordBatch.of(STREAM_3, 0, 10, 10, TestUtils.random(200)); objectWriter.write(STREAM_1, List.of(r1)); objectWriter.write(STREAM_3, List.of(r2, r3)); objectWriter.close().join(); @@ -618,6 +631,70 @@ public void testCompactNoneExistObjects2() { } } + @Test + public void testCompactWithUploadException() { + when(config.streamSetObjectCompactionStreamSplitSize()).thenReturn(100 * 1024 * 1024L); + when(config.streamSetObjectCompactionCacheSize()).thenReturn(1024 * 1024 * 1024L); + when(config.objectPartSize()).thenReturn(100 * 1024 * 1024); + Map> streamDataBlockMap = getStreamDataBlockMapLarge(); + S3ObjectMetadata objectMetadata0 = new S3ObjectMetadata(OBJECT_0, 0, S3ObjectType.STREAM_SET); + S3ObjectMetadata objectMetadata1 = new S3ObjectMetadata(OBJECT_1, 0, S3ObjectType.STREAM_SET); + S3ObjectMetadata objectMetadata2 = new S3ObjectMetadata(OBJECT_2, 0, S3ObjectType.STREAM_SET); + List s3ObjectMetadata = List.of(objectMetadata0, objectMetadata1, objectMetadata2); + this.compactionAnalyzer = new CompactionAnalyzer(config.streamSetObjectCompactionCacheSize(), config.streamSetObjectCompactionStreamSplitSize(), + config.maxStreamNumPerStreamSetObject(), config.maxStreamObjectNumPerCommit()); + List compactionPlans = this.compactionAnalyzer.analyze(streamDataBlockMap, new HashSet<>()); + CommitStreamSetObjectRequest request = new CommitStreamSetObjectRequest(); + + S3AsyncClient s3AsyncClient = Mockito.mock(S3AsyncClient.class); + doAnswer(invocation -> CompletableFuture.failedFuture(S3Exception.builder().statusCode(HttpStatusCode.NOT_FOUND).build())).when(s3AsyncClient).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); + doAnswer(invocation -> CompletableFuture.completedFuture(CreateMultipartUploadResponse.builder().uploadId("123").build())).when(s3AsyncClient).createMultipartUpload(any(CreateMultipartUploadRequest.class)); + doAnswer(invocation -> CompletableFuture.failedFuture(S3Exception.builder().statusCode(HttpStatusCode.NOT_FOUND).build())).when(s3AsyncClient).uploadPart(any(UploadPartRequest.class), any(AsyncRequestBody.class)); + + AwsObjectStorage objectStorage = Mockito.spy(new AwsObjectStorage(s3AsyncClient, "")); + doAnswer(invocation -> CompletableFuture.completedFuture(TestUtils.randomPooled(65 * 1024 * 1024))).when(objectStorage).rangeRead(any(), eq(objectMetadata0.key()), anyLong(), anyLong()); + doAnswer(invocation -> CompletableFuture.completedFuture(TestUtils.randomPooled(80 * 1024 * 1024))).when(objectStorage).rangeRead(any(), eq(objectMetadata1.key()), anyLong(), anyLong()); + doAnswer(invocation -> CompletableFuture.completedFuture(TestUtils.randomPooled(50 * 1024 * 1024))).when(objectStorage).rangeRead(any(), eq(objectMetadata2.key()), anyLong(), anyLong()); + + CompactionManager compactionManager = new CompactionManager(config, objectManager, streamManager, objectStorage); + Assertions.assertThrowsExactly(CompletionException.class, + () -> compactionManager.executeCompactionPlans(request, compactionPlans, s3ObjectMetadata)); + for (CompactionPlan plan : compactionPlans) { + plan.streamDataBlocksMap().forEach((streamId, blocks) -> blocks.forEach(block -> { + if (block.getObjectId() != OBJECT_1) { + block.getDataCf().thenAccept(data -> { + Assertions.assertEquals(0, data.refCnt()); + }).join(); + } + })); + } + } + + private static Map> getStreamDataBlockMapLarge() { + StreamDataBlock block1 = new StreamDataBlock(OBJECT_0, new DataBlockIndex(0, 0, 15, 15, 0, 15 * 1024 * 1024)); + StreamDataBlock block2 = new StreamDataBlock(OBJECT_0, new DataBlockIndex(1, 0, 20, 20, 15, 50 * 1024 * 1024)); + + StreamDataBlock block3 = new StreamDataBlock(OBJECT_1, new DataBlockIndex(0, 15, 12, 12, 0, 20 * 1024 * 1024)); + StreamDataBlock block4 = new StreamDataBlock(OBJECT_1, new DataBlockIndex(1, 20, 25, 25, 20, 60 * 1024 * 1024)); + + StreamDataBlock block5 = new StreamDataBlock(OBJECT_2, new DataBlockIndex(0, 27, 13, 20, 0, 20 * 1024 * 1024)); + StreamDataBlock block6 = new StreamDataBlock(OBJECT_2, new DataBlockIndex(3, 0, 30, 30, 20, 30 * 1024 * 1024)); + return Map.of( + OBJECT_0, List.of( + block1, + block2 + ), + OBJECT_1, List.of( + block3, + block4 + ), + OBJECT_2, List.of( + block5, + block6 + ) + ); + } + @Test public void testCompactWithLimit() { when(config.streamSetObjectCompactionStreamSplitSize()).thenReturn(70L); @@ -664,7 +741,7 @@ public void testCompactionShutdown() throws Throwable { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_0, objectId); ObjectWriter objectWriter = ObjectWriter.writer(objectId, objectStorage, 1024, 1024); - StreamRecordBatch r1 = new StreamRecordBatch(STREAM_0, 0, 0, 80, TestUtils.random(80)); + StreamRecordBatch r1 = StreamRecordBatch.of(STREAM_0, 0, 0, 80, TestUtils.random(80)); objectWriter.write(STREAM_0, List.of(r1)); objectWriter.close().join(); List streamsIndices = List.of( @@ -680,7 +757,7 @@ public void testCompactionShutdown() throws Throwable { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_1, objectId); ObjectWriter objectWriter = ObjectWriter.writer(OBJECT_1, objectStorage, 1024, 1024); - StreamRecordBatch r2 = new StreamRecordBatch(STREAM_0, 0, 80, 120, TestUtils.random(120)); + StreamRecordBatch r2 = StreamRecordBatch.of(STREAM_0, 0, 80, 120, TestUtils.random(120)); objectWriter.write(STREAM_0, List.of(r2)); objectWriter.close().join(); List streamsIndices = List.of( diff --git a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionTestBase.java b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionTestBase.java index 3af4d12b77..abb197ec0a 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionTestBase.java +++ b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionTestBase.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; @@ -79,10 +87,10 @@ public void setUp() throws Exception { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_0, objectId); ObjectWriter objectWriter = ObjectWriter.writer(objectId, objectStorage, 1024, 1024); - StreamRecordBatch r1 = new StreamRecordBatch(STREAM_0, 0, 0, 15, TestUtils.random(2)); - StreamRecordBatch r2 = new StreamRecordBatch(STREAM_1, 0, 25, 5, TestUtils.random(2)); - StreamRecordBatch r3 = new StreamRecordBatch(STREAM_1, 0, 30, 30, TestUtils.random(22)); - StreamRecordBatch r4 = new StreamRecordBatch(STREAM_2, 0, 30, 30, TestUtils.random(22)); + StreamRecordBatch r1 = StreamRecordBatch.of(STREAM_0, 0, 0, 15, TestUtils.random(2)); + StreamRecordBatch r2 = StreamRecordBatch.of(STREAM_1, 0, 25, 5, TestUtils.random(2)); + StreamRecordBatch r3 = StreamRecordBatch.of(STREAM_1, 0, 30, 30, TestUtils.random(22)); + StreamRecordBatch r4 = StreamRecordBatch.of(STREAM_2, 0, 30, 30, TestUtils.random(22)); objectWriter.write(STREAM_0, List.of(r1)); objectWriter.write(STREAM_1, List.of(r2)); objectWriter.write(STREAM_1, List.of(r3)); @@ -104,8 +112,8 @@ public void setUp() throws Exception { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_1, objectId); ObjectWriter objectWriter = ObjectWriter.writer(OBJECT_1, objectStorage, 1024, 1024); - StreamRecordBatch r5 = new StreamRecordBatch(STREAM_0, 0, 15, 5, TestUtils.random(1)); - StreamRecordBatch r6 = new StreamRecordBatch(STREAM_1, 0, 60, 60, TestUtils.random(52)); + StreamRecordBatch r5 = StreamRecordBatch.of(STREAM_0, 0, 15, 5, TestUtils.random(1)); + StreamRecordBatch r6 = StreamRecordBatch.of(STREAM_1, 0, 60, 60, TestUtils.random(52)); objectWriter.write(STREAM_0, List.of(r5)); objectWriter.write(STREAM_1, List.of(r6)); objectWriter.close().join(); @@ -123,8 +131,8 @@ public void setUp() throws Exception { objectManager.prepareObject(1, TimeUnit.MINUTES.toMillis(30)).thenAccept(objectId -> { assertEquals(OBJECT_2, objectId); ObjectWriter objectWriter = ObjectWriter.writer(OBJECT_2, objectStorage, 1024, 1024); - StreamRecordBatch r8 = new StreamRecordBatch(STREAM_1, 0, 400, 100, TestUtils.random(92)); - StreamRecordBatch r9 = new StreamRecordBatch(STREAM_2, 0, 230, 40, TestUtils.random(32)); + StreamRecordBatch r8 = StreamRecordBatch.of(STREAM_1, 0, 400, 100, TestUtils.random(92)); + StreamRecordBatch r9 = StreamRecordBatch.of(STREAM_2, 0, 230, 40, TestUtils.random(32)); objectWriter.write(STREAM_1, List.of(r8)); objectWriter.write(STREAM_2, List.of(r9)); objectWriter.close().join(); diff --git a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUploaderTest.java b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUploaderTest.java index 7d885f151a..27b209e411 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUploaderTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUploaderTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUtilTest.java b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUtilTest.java index a713d9d0a4..9e6bad9262 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUtilTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/compact/CompactionUtilTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; diff --git a/s3stream/src/test/java/com/automq/stream/s3/compact/StreamObjectCompactorTest.java b/s3stream/src/test/java/com/automq/stream/s3/compact/StreamObjectCompactorTest.java index 5757db2e1b..17421de1f0 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/compact/StreamObjectCompactorTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/compact/StreamObjectCompactorTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.compact; @@ -560,6 +568,6 @@ public void testCompactByCompositeObject() throws ExecutionException, Interrupte } StreamRecordBatch newRecord(long offset, int count, int payloadSize) { - return new StreamRecordBatch(streamId, 0, offset, count, TestUtils.random(payloadSize)); + return StreamRecordBatch.of(streamId, 0, offset, count, TestUtils.random(payloadSize)); } } diff --git a/s3stream/src/test/java/com/automq/stream/s3/failover/FailoverTest.java b/s3stream/src/test/java/com/automq/stream/s3/failover/FailoverTest.java deleted file mode 100644 index fc4bb1931a..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/failover/FailoverTest.java +++ /dev/null @@ -1,85 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.failover; - -import com.automq.stream.s3.wal.impl.block.BlockWALService; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.when; - -public class FailoverTest { - String path; - FailoverFactory failoverFactory; - WALRecover walRecover; - Failover failover; - - @BeforeEach - public void setup() { - path = "/tmp/" + System.currentTimeMillis() + "/failover_test_wal"; - failoverFactory = mock(FailoverFactory.class); - walRecover = mock(WALRecover.class); - failover = spy(new Failover(failoverFactory, walRecover)); - } - - @AfterEach - public void cleanup() throws IOException { - Files.delete(Path.of(path)); - } - - @Test - public void test() throws IOException, ExecutionException, InterruptedException, TimeoutException { - BlockWALService wal = BlockWALService.builder(path, 1024 * 1024).nodeId(233).epoch(100).build(); - wal.start(); - wal.shutdownGracefully(); - - FailoverRequest request = new FailoverRequest(); - - // node mismatch - request.setNodeId(234); - request.setDevice(path); - request.setVolumeId("test_volume_id"); - - when(failoverFactory.getWal(any())).thenAnswer(s -> - BlockWALService.builder(path, 1024 * 1024).nodeId(233).epoch(100).build()); - - boolean exceptionThrown = false; - try { - failover.failover(request).get(100, TimeUnit.SECONDS); - } catch (ExecutionException e) { - if (e.getCause() instanceof IllegalArgumentException) { - exceptionThrown = true; - } - } - Assertions.assertTrue(exceptionThrown); - - // node match - request.setNodeId(233); - FailoverResponse resp = failover.failover(request).get(1, TimeUnit.SECONDS); - assertEquals(233, resp.getNodeId()); - } - -} diff --git a/s3stream/src/test/java/com/automq/stream/s3/index/LocalStreamRangeIndexCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/index/LocalStreamRangeIndexCacheTest.java index b5d9cd46bd..d659259a94 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/index/LocalStreamRangeIndexCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/index/LocalStreamRangeIndexCacheTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; @@ -36,7 +44,7 @@ public class LocalStreamRangeIndexCacheTest { public void testInit() { ObjectStorage objectStorage = new MemoryObjectStorage(); // init with empty index - LocalStreamRangeIndexCache cache = new LocalStreamRangeIndexCache(); + LocalStreamRangeIndexCache cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); Assertions.assertEquals(-1, cache.searchObjectId(STREAM_0, 0).join()); @@ -47,7 +55,7 @@ public void testInit() { cache.updateIndexFromRequest(request); cache.upload().join(); - cache = new LocalStreamRangeIndexCache(); + cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); cache.initCf().join(); @@ -61,7 +69,7 @@ public void testInit() { @Test public void testAppend() { ObjectStorage objectStorage = new MemoryObjectStorage(); - LocalStreamRangeIndexCache cache = new LocalStreamRangeIndexCache(); + LocalStreamRangeIndexCache cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); CommitStreamSetObjectRequest request = new CommitStreamSetObjectRequest(); @@ -85,7 +93,7 @@ public void testAppend() { @Test public void testPrune() { ObjectStorage objectStorage = new MemoryObjectStorage(); - LocalStreamRangeIndexCache cache = new LocalStreamRangeIndexCache(); + LocalStreamRangeIndexCache cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); CommitStreamSetObjectRequest request = new CommitStreamSetObjectRequest(); @@ -122,7 +130,7 @@ public void testPrune() { Assertions.assertEquals(97, cache.searchObjectId(STREAM_0, 1500).join()); // test load from object storage - cache = new LocalStreamRangeIndexCache(); + cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); cache.initCf().join(); @@ -147,7 +155,7 @@ public void testPrune() { @Test public void testEvict() { ObjectStorage objectStorage = new MemoryObjectStorage(); - LocalStreamRangeIndexCache cache = new LocalStreamRangeIndexCache(); + LocalStreamRangeIndexCache cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); int streamNum = 500; @@ -172,7 +180,7 @@ public void testEvict() { @Test public void testCompact() { ObjectStorage objectStorage = new MemoryObjectStorage(); - LocalStreamRangeIndexCache cache = new LocalStreamRangeIndexCache(); + LocalStreamRangeIndexCache cache = LocalStreamRangeIndexCache.create(); cache.start(); cache.init(NODE_0, objectStorage); CommitStreamSetObjectRequest request = new CommitStreamSetObjectRequest(); @@ -249,6 +257,42 @@ public void testCompact() { Assertions.assertEquals(-1, cache.searchObjectId(STREAM_0, 300).join()); } + @Test + public void testCompactWithStreamDeleted() { + ObjectStorage objectStorage = new MemoryObjectStorage(); + LocalStreamRangeIndexCache cache = LocalStreamRangeIndexCache.create(); + cache.start(); + cache.init(NODE_0, objectStorage); + CommitStreamSetObjectRequest request = new CommitStreamSetObjectRequest(); + long startOffset = 50; + for (int i = 0; i < 10; i++) { + request.setObjectId(88 + i); + request.setStreamRanges(List.of( + new ObjectStreamRange(STREAM_0, 0, startOffset, startOffset + 100, 100), + new ObjectStreamRange(STREAM_1, 0, startOffset, startOffset + 100, 100))); + cache.updateIndexFromRequest(request).join(); + startOffset += 100; + } + Assertions.assertEquals(10, cache.getStreamRangeIndexMap().get(STREAM_0).length()); + Assertions.assertEquals(10, cache.getStreamRangeIndexMap().get(STREAM_1).length()); + Assertions.assertEquals(20 * RangeIndex.OBJECT_SIZE, cache.totalSize()); + + // mock STREAM_0 deleted + request.setObjectId(256); + request.setStreamRanges(List.of( + new ObjectStreamRange(STREAM_1, 0, 50, 1050, 1000) + )); + request.setCompactedObjectIds(List.of(88L, 89L, 90L, 91L, 92L, 93L, 94L, 95L, 96L, 97L)); + request.setStreamObjects(Collections.emptyList()); + cache.updateIndexFromRequest(request).join(); + + Assertions.assertNull(cache.getStreamRangeIndexMap().get(STREAM_0)); + Assertions.assertEquals(1, cache.getStreamRangeIndexMap().get(STREAM_1).getRangeIndexList().size()); + Assertions.assertEquals(RangeIndex.OBJECT_SIZE, cache.totalSize()); + Assertions.assertEquals(new RangeIndex(50, 1050, 256), + cache.getStreamRangeIndexMap().get(STREAM_1).getRangeIndexList().get(0)); + } + private StreamObject newStreamObject(long objectId, long objectSize, long streamId, long startOffset, long endOffset) { StreamObject streamObject = new StreamObject(); streamObject.setObjectId(objectId); diff --git a/s3stream/src/test/java/com/automq/stream/s3/index/MockRandom.java b/s3stream/src/test/java/com/automq/stream/s3/index/MockRandom.java index f0c1134ed2..38727c1abe 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/index/MockRandom.java +++ b/s3stream/src/test/java/com/automq/stream/s3/index/MockRandom.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; diff --git a/s3stream/src/test/java/com/automq/stream/s3/index/NodeRangeIndexCacheTest.java b/s3stream/src/test/java/com/automq/stream/s3/index/NodeRangeIndexCacheTest.java index ef419685bb..87aacb49be 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/index/NodeRangeIndexCacheTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/index/NodeRangeIndexCacheTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; diff --git a/s3stream/src/test/java/com/automq/stream/s3/index/SparseRangeIndexTest.java b/s3stream/src/test/java/com/automq/stream/s3/index/SparseRangeIndexTest.java index 02241bb4c6..bcde500002 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/index/SparseRangeIndexTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/index/SparseRangeIndexTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.index; diff --git a/s3stream/src/test/java/com/automq/stream/s3/metrics/AttributesUtilTest.java b/s3stream/src/test/java/com/automq/stream/s3/metrics/AttributesUtilTest.java index 3ff7b81b8e..b087d13c2f 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/metrics/AttributesUtilTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/metrics/AttributesUtilTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.metrics; diff --git a/s3stream/src/test/java/com/automq/stream/s3/model/StreamRecordBatchTest.java b/s3stream/src/test/java/com/automq/stream/s3/model/StreamRecordBatchTest.java new file mode 100644 index 0000000000..cd58659f6b --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/model/StreamRecordBatchTest.java @@ -0,0 +1,81 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.model; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.nio.charset.StandardCharsets; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.buffer.Unpooled; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class StreamRecordBatchTest { + + @Test + public void testOf() { + byte[] payload = "hello".getBytes(StandardCharsets.UTF_8); + ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload); + StreamRecordBatch record = StreamRecordBatch.of(1L, 2L, 3L, 4, payloadBuf); + assertEquals(1, record.getStreamId()); + assertEquals(2, record.getEpoch()); + assertEquals(3, record.getBaseOffset()); + assertEquals(4, record.getCount()); + assertEquals(payload.length, record.size()); + assertEquals(0, payloadBuf.refCnt()); + byte[] realPayload = new byte[payload.length]; + record.getPayload().readBytes(realPayload); + assertArrayEquals(payload, realPayload); + record.release(); + assertEquals(0, record.encoded.refCnt()); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + public void testParse(boolean duplicated) { + CompositeByteBuf buf = Unpooled.compositeBuffer(); + for (int i = 0; i < 10; i++) { + ByteBuf payloadBuf = Unpooled.wrappedBuffer(("hello" + i).getBytes(StandardCharsets.UTF_8)); + StreamRecordBatch record = StreamRecordBatch.of(1L, 2L, 3L + i, 4, payloadBuf); + buf.addComponent(true, record.encoded()); + } + for (int i = 0; i < 10; i++) { + StreamRecordBatch record = StreamRecordBatch.parse(buf, duplicated); + assertEquals(3 + i, record.getBaseOffset()); + ByteBuf payloadBuf = record.getPayload(); + byte[] payload = new byte[payloadBuf.readableBytes()]; + payloadBuf.readBytes(payload); + assertArrayEquals(("hello" + i).getBytes(StandardCharsets.UTF_8), payload); + record.release(); + if (duplicated) { + assertEquals(0, record.encoded.refCnt()); + } + } + assertEquals(0, buf.readableBytes()); + assertEquals(1, buf.refCnt()); + buf.release(); + } + +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/objects/ObjectManagerTest.java b/s3stream/src/test/java/com/automq/stream/s3/objects/ObjectManagerTest.java index 7875c87dc3..4feaaf3c1d 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/objects/ObjectManagerTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/objects/ObjectManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.objects; diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/AbstractObjectStorageTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/AbstractObjectStorageTest.java index 7cbbbcd813..8169151c48 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/operator/AbstractObjectStorageTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/AbstractObjectStorageTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -21,25 +29,37 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; import io.netty.buffer.ByteBuf; +import static org.awaitility.Awaitility.await; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.anyLong; import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.timeout; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @Tag("S3Unit") +@Timeout(10) class AbstractObjectStorageTest { AbstractObjectStorage objectStorage; @@ -128,6 +148,19 @@ void testMergeRead() throws ExecutionException, InterruptedException { buf.release(); } + @Test + void testHandleReadCompleted() throws Throwable { + ByteBuf data = TestUtils.random(4096); + CompletableFuture readToEndCf = new CompletableFuture<>(); + CompletableFuture readRangeCf = new CompletableFuture<>(); + AbstractObjectStorage.MergedReadTask.handleReadCompleted(List.of( + new AbstractObjectStorage.ReadTask(new ReadOptions(), "fake", 3000, -1, readToEndCf), + new AbstractObjectStorage.ReadTask(new ReadOptions(), "fake", 2000, 4096, readRangeCf) + ), 2000, data.slice(2000, 4096 - 2000), null); + assertEquals(data.slice(3000, 4096 - 3000), readToEndCf.get()); + assertEquals(data.slice(2000, 4096 - 2000), readRangeCf.get()); + } + @Test void testByteBufRefCnt() throws ExecutionException, InterruptedException { objectStorage = new MemoryObjectStorage(false); @@ -142,6 +175,198 @@ void testByteBufRefCnt() throws ExecutionException, InterruptedException { }).get(); } + @Test + void testFastRetry() throws Throwable { + // Initialize memory storage and spy to track method calls + objectStorage = new MemoryObjectStorage(); + objectStorage = spy(objectStorage); + + // Configure write options: enable fast retry, disable normal retry + ObjectStorage.WriteOptions options = new ObjectStorage.WriteOptions() + .enableFastRetry(true) + .retry(false); + + // Mock S3 latency calculator via reflection to force fast retry condition + Field latencyCalculatorField = AbstractObjectStorage.class.getDeclaredField("s3LatencyCalculator"); + latencyCalculatorField.setAccessible(true); + S3LatencyCalculator mockCalculator = mock(S3LatencyCalculator.class); + when(mockCalculator.valueAtPercentile(anyLong(), anyLong())).thenReturn(100L); // Force low latency to trigger fast retry + latencyCalculatorField.set(objectStorage, mockCalculator); + + // Track doWrite() calls: first call hangs, second completes immediately + AtomicInteger callCount = new AtomicInteger(); + CompletableFuture firstFuture = new CompletableFuture<>(); + when(objectStorage.doWrite(any(), anyString(), any())).thenAnswer(inv -> { + int count = callCount.getAndIncrement(); + return (count == 0) ? firstFuture : CompletableFuture.completedFuture(null); + }); + + // Execute write operation + ByteBuf data = TestUtils.randomPooled(1024); + assertEquals(1, data.refCnt()); // Verify initial ref count + + CompletableFuture writeFuture = objectStorage.write(options, "testKey", data); + writeFuture.get(1, TimeUnit.SECONDS); // Wait for write completion + + // Verify: two calls made (initial + retry), data ref count maintained during retry + assertEquals(1, data.refCnt()); + assertEquals(2, callCount.get()); + + // Complete initial future and verify data release + firstFuture.complete(null); + await().atMost(1, TimeUnit.SECONDS) + .untilAsserted(() -> assertEquals(0, data.refCnt())); // Ensure buffer released + } + + @Test + void testWriteRetryTimeout() throws Throwable { + // Setup storage with 100ms timeout (clearer time unit) + objectStorage = spy(new MemoryObjectStorage()); + ObjectStorage.WriteOptions options = new ObjectStorage.WriteOptions() + .retry(true) + .timeout(1000L); + + // Mock hanging write operation + AtomicInteger callCount = new AtomicInteger(); + when(objectStorage.doWrite(any(), anyString(), any())).thenAnswer(inv -> { + int count = callCount.getAndIncrement(); + if (count < 12) { + CompletableFuture future = new CompletableFuture<>(); + Executors.newSingleThreadScheduledExecutor().schedule( + () -> future.completeExceptionally(new TimeoutException("Simulated timeout")), + 100, TimeUnit.MILLISECONDS + ); + return future; + } + // Second call: immediate success + return CompletableFuture.completedFuture(null); + }); + + // Execute test + ByteBuf data = TestUtils.randomPooled(1024); + CompletableFuture writeFuture = + objectStorage.write(options, "testKey", data); + // Verify timeout exception + assertThrows(TimeoutException.class, + () -> writeFuture.get(1, TimeUnit.SECONDS)); + // Verify resource cleanup + await().atMost(2, TimeUnit.SECONDS) + .untilAsserted(() -> assertEquals(0, data.refCnt())); + // Verify: no successful calls made + assertTrue(callCount.get() < 12); + } + + @Test + void testWritePermit() throws Exception { + final int maxConcurrency = 5; + objectStorage = spy(new MemoryObjectStorage(maxConcurrency)); + + ObjectStorage.WriteOptions options = new ObjectStorage.WriteOptions() + .enableFastRetry(false) + .retry(false); + + // Use completable future to block first 5 calls + CompletableFuture barrierFuture = new CompletableFuture<>(); + AtomicInteger callCount = new AtomicInteger(); + + when(objectStorage.doWrite(any(), anyString(), any())).thenAnswer(inv -> { + int count = callCount.getAndIncrement(); + return (count < maxConcurrency) + ? barrierFuture // Block first 5 calls + : CompletableFuture.completedFuture(null); // Immediate success for 6th + }); + + // Phase 1: Submit max concurrency requests + List buffers = new ArrayList<>(); + for (int i = 0; i < maxConcurrency; i++) { + ByteBuf data = TestUtils.randomPooled(1024); + buffers.add(data); + objectStorage.write(options, "testKey", data); + } + + // Verify initial calls reached max concurrency + await().atMost(1, TimeUnit.SECONDS) + .untilAsserted(() -> assertEquals(maxConcurrency, callCount.get())); + + // Phase 2: Submit 6th request beyond concurrency limit + CompletableFuture sixthWriteFuture = + CompletableFuture.supplyAsync(() -> + objectStorage.write(options, "testKey", TestUtils.random(1024)) + ).thenCompose(f -> f); + + // Release blocked calls and verify completion + barrierFuture.complete(null); + await().atMost(2, TimeUnit.SECONDS) + .untilAsserted(() -> { + assertEquals(maxConcurrency + 1, callCount.get()); + assertTrue(sixthWriteFuture.isDone()); + + // Verify: all buffers released + for (ByteBuf buffer : buffers) { + assertEquals(0, buffer.refCnt()); + } + }); + } + + @Test + void testWaitWritePermit() throws Exception { + final int maxConcurrency = 1; + objectStorage = spy(new MemoryObjectStorage(maxConcurrency)); + + ObjectStorage.WriteOptions options = new ObjectStorage.WriteOptions() + .enableFastRetry(false) + .retry(false); + + // Block first call using completable future + CompletableFuture blockingFuture = new CompletableFuture<>(); + AtomicInteger callCount = new AtomicInteger(); + + when(objectStorage.doWrite(any(), anyString(), any())).thenAnswer(inv -> { + callCount.incrementAndGet(); + return blockingFuture; // Always return blocking future for first call + }); + + // Phase 1: Acquire the only permit + ByteBuf firstBuffer = TestUtils.randomPooled(1024); + objectStorage.write(options, "testKey", firstBuffer); + + // Verify permit acquisition + await().until(() -> callCount.get() == 1); + + // Phase 2: Verify blocking behavior with interrupt + Thread blockingThread = new Thread(() -> { + ByteBuf byteBuf = TestUtils.randomPooled(1024); + try { + CompletableFuture future = + objectStorage.write(options, "testKey", byteBuf); + ExecutionException exception = assertThrows(ExecutionException.class, () -> future.get()); + assertTrue(exception.getCause() instanceof InterruptedException); + } catch (Exception e) { + // Ignore + } finally { + await().atMost(1, TimeUnit.SECONDS).untilAsserted(() -> { + assertEquals(0, byteBuf.refCnt()); + }); + } + }); + + blockingThread.start(); + + Thread.sleep(1000); + + // Interrupt and verify + blockingThread.interrupt(); + blockingThread.join(); + + // Verify resource cleanup + assertEquals(1, firstBuffer.refCnt()); + + // Cleanup + blockingFuture.complete(null); + await().atMost(2, TimeUnit.SECONDS) + .untilAsserted(() -> assertEquals(0, firstBuffer.refCnt())); + } + @Test void testReadToEndOfObject() throws ExecutionException, InterruptedException { objectStorage = new MemoryObjectStorage(true); diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/AwsObjectStorageTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/AwsObjectStorageTest.java new file mode 100644 index 0000000000..52c4cfad97 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/AwsObjectStorageTest.java @@ -0,0 +1,43 @@ +package com.automq.stream.s3.operator; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.core.SdkSystemSetting; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doCallRealMethod; +import static org.mockito.Mockito.mock; + +public class AwsObjectStorageTest { + + @Test + void testCredentialsProviderChain() { + AwsObjectStorage storage = mock(AwsObjectStorage.class); + doCallRealMethod().when(storage).credentialsProviders0(any()); + doCallRealMethod().when(storage).newCredentialsProviderChain(any()); + + AwsCredentialsProvider provider = storage.newCredentialsProviderChain(storage.credentialsProviders0( + BucketURI.parse("0@s3://bucket?region=us-east-1&accessKey=ak&secretKey=sk"))); + AwsCredentials credentials = provider.resolveCredentials(); + Assertions.assertInstanceOf(AwsBasicCredentials.class, credentials); + AwsBasicCredentials basicCredentials = (AwsBasicCredentials) credentials; + Assertions.assertEquals("ak", basicCredentials.accessKeyId()); + Assertions.assertEquals("sk", basicCredentials.secretAccessKey()); + + // test fallback to system property credential provider + + System.setProperty(SdkSystemSetting.AWS_ACCESS_KEY_ID.property(), "ak"); + System.setProperty(SdkSystemSetting.AWS_SECRET_ACCESS_KEY.property(), "sk"); + provider = storage.newCredentialsProviderChain(storage.credentialsProviders0( + BucketURI.parse("0@s3://bucket?region=us-east-1&accessKey=&secretKey="))); + credentials = provider.resolveCredentials(); + Assertions.assertInstanceOf(AwsBasicCredentials.class, credentials); + basicCredentials = (AwsBasicCredentials) credentials; + Assertions.assertEquals("ak", basicCredentials.accessKeyId()); + Assertions.assertEquals("sk", basicCredentials.secretAccessKey()); + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/BucketURITest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/BucketURITest.java index 7321de194e..a8a896b610 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/operator/BucketURITest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/BucketURITest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -27,7 +35,7 @@ public class BucketURITest { public void testParse_valid() { String bucketStr = "0@s3://bucket1?region=region1&k1=v1&k2=v2&k2=v22&endpoint=https://aws.amazon.com:444," + "1@gs://bucket2?region=region2&endpoint=https://gcp," + - "2@azblob://bucket3"; + "-2@azblob://bucket3"; List buckets = BucketURI.parseBuckets(bucketStr); assertEquals((short) 0, buckets.get(0).bucketId()); assertEquals("bucket1", buckets.get(0).bucket()); @@ -45,7 +53,7 @@ public void testParse_valid() { assertEquals("https://gcp", buckets.get(1).endpoint()); assertEquals("gs", buckets.get(1).protocol()); - assertEquals((short) 2, buckets.get(2).bucketId()); + assertEquals((short) -2, buckets.get(2).bucketId()); assertEquals("bucket3", buckets.get(2).bucket()); assertEquals("", buckets.get(2).region()); assertEquals("", buckets.get(2).endpoint()); diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/DeleteObjectsAccumulatorTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/DeleteObjectsAccumulatorTest.java index 6baacd753a..c283bf3469 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/operator/DeleteObjectsAccumulatorTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/DeleteObjectsAccumulatorTest.java @@ -1,17 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Threads; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; @@ -31,7 +40,6 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -39,6 +47,7 @@ import java.util.stream.Collectors; import java.util.stream.IntStream; +import static com.automq.stream.s3.operator.DeleteObjectsAccumulator.LOGGER; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -241,7 +250,7 @@ void testHighTrafficBatchDelete() { CountDownLatch latch = new CountDownLatch(batchNumber); - ExecutorService executorService = Executors.newFixedThreadPool(batchNumber); + ExecutorService executorService = Threads.newFixedFastThreadLocalThreadPoolWithMonitor(batchNumber, "delete-obj-accumulator-thread", true, LOGGER); for (int j = 0; j < batchNumber; j++) { int finalJ = j; executorService.submit(() -> { diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/LocalFileObjectStorageTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/LocalFileObjectStorageTest.java new file mode 100644 index 0000000000..9b5a979cbf --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/LocalFileObjectStorageTest.java @@ -0,0 +1,155 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.exceptions.ObjectNotExistException; +import com.automq.stream.s3.metadata.ObjectUtils; +import com.automq.stream.utils.FutureUtil; +import com.automq.stream.utils.Utils; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; + +@Tag("S3Unit") +public class LocalFileObjectStorageTest { + + LocalFileObjectStorage objectStorage; + String base = "/tmp/automq_test/localfilestoragetest"; + + @BeforeEach + public void setup() { + objectStorage = new LocalFileObjectStorage(BucketURI.parse("-2@file://" + base)); + } + + @AfterEach + public void cleanup() throws IOException { + Utils.delete(Path.of(base)); + } + + @Test + public void testReadWrite() throws ExecutionException, InterruptedException { + String key = ObjectUtils.genKey(0, 100); + Writer writer = objectStorage.writer(new ObjectStorage.WriteOptions(), key); + writer.write(Unpooled.wrappedBuffer("hello ".getBytes(StandardCharsets.UTF_8))); + byte[] bytes = "world".getBytes(StandardCharsets.UTF_8); + ByteBuf buf = ByteBufAlloc.byteBuffer(bytes.length); + buf.writeBytes(bytes); + writer.write(buf); + writer.close().get(); + assertEquals(0, buf.refCnt()); + + buf = objectStorage.rangeRead(new ObjectStorage.ReadOptions(), key, 0, -1L).get(); + assertEquals("hello world", substr(buf, 0, buf.readableBytes())); + assertEquals("hello", substr(buf, 0, 5)); + + objectStorage.delete(List.of(new ObjectStorage.ObjectInfo(objectStorage.bucketId(), key, 0, 0))).get(); + + Throwable exception = null; + try { + objectStorage.rangeRead(new ObjectStorage.ReadOptions(), key, 0, -1L).get(); + } catch (Throwable e) { + exception = FutureUtil.cause(e); + } + assertEquals(ObjectNotExistException.class, Optional.ofNullable(exception).map(Throwable::getClass).orElse(null)); + } + + @Test + public void testList() throws ExecutionException, InterruptedException { + objectStorage.write(new ObjectStorage.WriteOptions(), "abc/def/100", Unpooled.wrappedBuffer("hello world".getBytes(StandardCharsets.UTF_8))).get(); + objectStorage.write(new ObjectStorage.WriteOptions(), "abc/def/101", Unpooled.wrappedBuffer("hello world1".getBytes(StandardCharsets.UTF_8))).get(); + objectStorage.write(new ObjectStorage.WriteOptions(), "abc/deg/102", Unpooled.wrappedBuffer("hello world2".getBytes(StandardCharsets.UTF_8))).get(); + + assertEquals( + List.of("abc/def/100", "abc/def/101", "abc/deg/102"), + objectStorage.list("").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + assertEquals( + List.of("abc/def/100", "abc/def/101"), + objectStorage.list("abc/def").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + assertEquals( + List.of("abc/def/100", "abc/def/101", "abc/deg/102"), + objectStorage.list("abc/de").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + assertEquals( + List.of("abc/def/100", "abc/def/101", "abc/deg/102"), + objectStorage.list("ab").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + assertEquals( + List.of("abc/def/100", "abc/def/101"), + objectStorage.list("abc/def/").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + assertEquals( + List.of("abc/def/100", "abc/def/101"), + objectStorage.list("abc/def/1").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + assertEquals( + Collections.emptyList(), + objectStorage.list("abc/deh").get().stream().map(ObjectStorage.ObjectPath::key).sorted().collect(Collectors.toList()) + ); + } + + @Test + public void testDiskFull() throws Throwable { + objectStorage.availableSpace.set(10); + String key = ObjectUtils.genKey(0, 100); + objectStorage.write(new ObjectStorage.WriteOptions(), "abc/def/100", Unpooled.wrappedBuffer("hhhhhhhhh".getBytes(StandardCharsets.UTF_8))).get(); + CompletableFuture w2 = objectStorage.write(new ObjectStorage.WriteOptions(), "abc/def/101", Unpooled.wrappedBuffer("h2".getBytes(StandardCharsets.UTF_8))); + CompletableFuture w3 = objectStorage.write(new ObjectStorage.WriteOptions(), "abc/def/102", Unpooled.wrappedBuffer("h3".getBytes(StandardCharsets.UTF_8))); + assertEquals(2, objectStorage.waitingTasks.size()); + assertEquals(1, objectStorage.availableSpace.get()); + assertFalse(w2.isDone()); + assertFalse(w3.isDone()); + objectStorage.delete(List.of(new ObjectStorage.ObjectInfo(objectStorage.bucketId(), "abc/def/100", 0, 0))).get(); + w2.get(1, TimeUnit.SECONDS); + w3.get(1, TimeUnit.SECONDS); + assertEquals(0, objectStorage.waitingTasks.size()); + assertEquals(6, objectStorage.availableSpace.get()); + } + + private String substr(ByteBuf buf, int start, int end) { + buf = buf.duplicate(); + byte[] bytes = new byte[end - start]; + buf.skipBytes(start); + buf.readBytes(bytes); + return new String(bytes, StandardCharsets.UTF_8); + } + +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/MultiPartWriterTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/MultiPartWriterTest.java index 660d9ebb33..f7c3797986 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/operator/MultiPartWriterTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/MultiPartWriterTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -68,7 +76,7 @@ void setUp() { @Test void testWrite() throws NoSuchMethodException, InvocationTargetException, IllegalAccessException, ExecutionException, InterruptedException { - writer = new MultiPartWriter(ObjectStorage.WriteOptions.DEFAULT, operator, "test-path", 100); + writer = new MultiPartWriter(ObjectStorage.WriteOptions.DEFAULT, operator, "test-path", 100, 100); List requests = new ArrayList<>(); List contentLengths = new ArrayList<>(); diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/ProxyWriterTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/ProxyWriterTest.java index 2c9c92476e..4a6cb790ae 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/operator/ProxyWriterTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/ProxyWriterTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; @@ -69,10 +77,10 @@ public void testWrite_dataLargerThanMaxUploadSize() { writer.write(TestUtils.random(17 * 1024 * 1024)); assertTrue(writer.hasBatchingPart()); assertNull(writer.largeObjectWriter); - writer.write(TestUtils.random(17 * 1024 * 1024)); + writer.write(TestUtils.random(33 * 1024 * 1024)); assertNotNull(writer.largeObjectWriter); assertFalse(writer.hasBatchingPart()); - writer.write(TestUtils.random(17 * 1024 * 1024)); + writer.write(TestUtils.random(33 * 1024 * 1024)); assertNotNull(writer.largeObjectWriter); assertFalse(writer.hasBatchingPart()); writer.close(); diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/S3LatencyCalculatorTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/S3LatencyCalculatorTest.java index fa84ae7d88..72067d6bcc 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/operator/S3LatencyCalculatorTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/S3LatencyCalculatorTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.operator; diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/TrafficRateLimiterTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/TrafficRateLimiterTest.java new file mode 100644 index 0000000000..026e94161b --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/TrafficRateLimiterTest.java @@ -0,0 +1,93 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; +import com.automq.stream.utils.ThreadUtils; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +@Tag("S3Unit") +public class TrafficRateLimiterTest { + + private ScheduledExecutorService scheduler; + + @BeforeEach + public void setUp() { + scheduler = Executors.newScheduledThreadPool(1); + } + + @AfterEach + public void tearDown() { + scheduler.shutdown(); + } + + @Test + public void testExceedsBoundary() { + TrafficRateLimiter limiter = new TrafficRateLimiter(scheduler); + long prev = limiter.currentRate(); + limiter.update(Long.MAX_VALUE); + assertEquals(prev, limiter.currentRate()); + limiter.update(0); + assertEquals(1L << 10, limiter.currentRate()); + } + + @Test + public void testConsumeBeforeUpdate() { + long rateLimit = 1024 * 1024; + long totalTraffic = 1024 * 1024 * 5; + ExecutorService executor = Executors.newSingleThreadExecutor(); + CountDownLatch consumeStarted = new CountDownLatch(1); + TrafficRateLimiter limiter = new TrafficRateLimiter(scheduler, rateLimit); + Future future = executor.submit(() -> { + long startTime = System.currentTimeMillis(); + limiter.consume(totalTraffic).join(); + consumeStarted.countDown(); + long endTime = System.currentTimeMillis(); + return endTime - startTime; + }); + + try { + consumeStarted.await(); // make sure update after the consume method is called + long prevRate = limiter.currentRate(); + limiter.update(0); + long duration = future.get(); + double actualRate = ((double) totalTraffic / 1024 / duration) * 1000; + assertTrue(actualRate > limiter.currentRate() && actualRate <= prevRate); + assertTrue(duration / 1000 <= 5); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } finally { + ThreadUtils.shutdownExecutor(executor, 1, TimeUnit.SECONDS); + } + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/TrafficRegulatorTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/TrafficRegulatorTest.java new file mode 100644 index 0000000000..7d46a0af64 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/TrafficRegulatorTest.java @@ -0,0 +1,270 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.operator; + +import com.google.common.collect.EvictingQueue; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; + +import java.lang.reflect.Field; + +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Tag("S3Unit") +class TrafficRegulatorTest { + + private static final long MIN_RATE_LIMITER_RATE = kbpsToBps(1); + private static final long MAX_RATE_LIMITER_RATE = TrafficRateLimiter.MAX_BUCKET_TOKENS_PER_SECOND << 10; + private static final long NEAR_MAX_RATE_LIMITER_RATE = (long) (MAX_RATE_LIMITER_RATE * 0.95); + + private static final long FAILURE_RATE = mbpsToBps(10); + + private TrafficMonitor successMonitor; + private TrafficMonitor failureMonitor; + private TrafficRateLimiter rateLimiter; + private TrafficVolumeLimiter volumeLimiter; + private Logger logger; + private TrafficRegulator regulator; + + /** + * Converts a rate given in MB/s to bytes/s. + */ + private static long mbpsToBps(double mbRate) { + return (long) (mbRate * (1 << 20)); + } + + /** + * Converts a rate given in KB/s to bytes/s. + */ + private static long kbpsToBps(double kbRate) { + return (long) (kbRate * (1 << 10)); + } + + @BeforeEach + void setUp() { + successMonitor = mock(TrafficMonitor.class); + failureMonitor = mock(TrafficMonitor.class); + rateLimiter = mock(TrafficRateLimiter.class); + volumeLimiter = mock(TrafficVolumeLimiter.class); + logger = mock(Logger.class); + regulator = new TrafficRegulator("testOperation", successMonitor, failureMonitor, rateLimiter, volumeLimiter, logger); + } + + // ---------------- Decrease tests (failure rate value does not affect if greater than 0) ---------------- + + @Test + void testRegulateDecreaseSuccessAboveMinWithFailure() { + long successRate = mbpsToBps(100); + setRegulatorDecreaseEnv(successRate, MIN_RATE_LIMITER_RATE); + regulator.regulate(); + checkRegulate(successRate); + } + + @Test + void testRegulateDecreaseSuccessBelowMinWithFailure() { + long successRate = mbpsToBps(5); + setRegulatorDecreaseEnv(successRate, MIN_RATE_LIMITER_RATE); + regulator.regulate(); + checkRegulate(getMinRateFromRegulator()); + } + + // ---------------- Increase tests ---------------- + + @Test + void testRegulateIncreaseWithMaxRateLimiter() { + setRegulatorIncreaseEnv(0, MAX_RATE_LIMITER_RATE); + regulator.regulate(); + checkRegulate(MAX_RATE_LIMITER_RATE); + } + + @Test + void testRegulateIncreaseWithSuccessBelowMinAndRateLimiterMinNoHistory() { + long successRate = mbpsToBps(5); + setRegulatorIncreaseEnv(successRate, MIN_RATE_LIMITER_RATE); + regulator.regulate(); + checkRegulate((long) (MIN_RATE_LIMITER_RATE + successRate * getFastIncrementRatio())); + } + + @Test + void testRegulateIncreaseWithSuccessBelowMinAndRateLimiterNearMaxNoHistory() { + long successRate = mbpsToBps(5); + setRegulatorIncreaseEnv(successRate, NEAR_MAX_RATE_LIMITER_RATE); + regulator.regulate(); + checkRegulate(MAX_RATE_LIMITER_RATE); + } + + @Test + void testRegulateIncreaseWithSuccessNearMaxAndRateLimiterNearMaxNoHistory() { + long successRate = (long) (getMaxRateFromRegulator() * 0.95); + setRegulatorIncreaseEnv(successRate, NEAR_MAX_RATE_LIMITER_RATE); + regulator.regulate(); + checkRegulate((long) (NEAR_MAX_RATE_LIMITER_RATE + successRate * getSlowIncrementRatio())); + } + + // ---------------- Tests involving success history ---------------- + + @Test + void testRegulateIncreaseWithHistoryNotFull() { + // Populate the success history queue with 5 entries + EvictingQueue queue = getSuccessRateQueue(); + queue.add((double) mbpsToBps(10.0)); + queue.add((double) mbpsToBps(20.0)); + queue.add((double) mbpsToBps(30.0)); + queue.add((double) mbpsToBps(40.0)); + queue.add((double) mbpsToBps(50.0)); + // Setup current rate to 60 MB/s, failure rate 0, success rate 0 + when(rateLimiter.currentRate()).thenReturn(mbpsToBps(60.0)); + when(successMonitor.getRateAndReset()).thenReturn(0.0); + when(failureMonitor.getRateAndReset()).thenReturn(0.0); + regulator.regulate(); + // Expected new rate: second largest of [60 + 50*0.5=85, 60 +50*0.05=62.5, 50] → 62.5 MB/s + long expectedNewRate = mbpsToBps(62.5); + verify(rateLimiter).update(expectedNewRate); + verify(volumeLimiter).update(expectedNewRate * getWindowSize()); + } + + @Test + void testRegulateIncreaseWithFullHistory() { + // Inject 64 entries into the queue (top 4: 100, 90, 80, 70) + EvictingQueue queue = getSuccessRateQueue(); + queue.add((double) mbpsToBps(100.0)); + queue.add((double) mbpsToBps(90.0)); + queue.add((double) mbpsToBps(80.0)); + queue.add((double) mbpsToBps(70.0)); + for (int i = 0; i < 60; i++) { + queue.add((double) mbpsToBps(60.0)); + } + // Setup current rate 80 MB/s, failure rate 0, success rate 0 + when(rateLimiter.currentRate()).thenReturn(mbpsToBps(80.0)); + when(successMonitor.getRateAndReset()).thenReturn(0.0); + when(failureMonitor.getRateAndReset()).thenReturn(0.0); + regulator.regulate(); + // Expected new rate: 85 MB/s (mean of top 4 entries) + long expectedNewRate = mbpsToBps(85.0); + verify(rateLimiter).update(expectedNewRate); + verify(volumeLimiter).update(expectedNewRate * getWindowSize()); + } + + @Test + void testRegulateIncreaseJumpsToMaxWhenCurrentLimitExceedsThreshold() { + // Setup history with mean 100 MB/s + EvictingQueue queue = getSuccessRateQueue(); + for (int i = 0; i < 4; i++) { + queue.add((double) mbpsToBps(100.0)); + } + for (int i = 0; i < 60; i++) { + queue.add((double) mbpsToBps(50.0)); + } + // Current rate is 701 MB/s (exceeds 7x history rate) + when(rateLimiter.currentRate()).thenReturn(mbpsToBps(701.0)); + when(successMonitor.getRateAndReset()).thenReturn(0.0); + when(failureMonitor.getRateAndReset()).thenReturn(0.0); + regulator.regulate(); + // Verify rate jumps to MAX + verify(rateLimiter).update(MAX_RATE_LIMITER_RATE); + verify(volumeLimiter).update(MAX_RATE_LIMITER_RATE * getWindowSize()); + } + + @Test + void testRegulateFailureDoesNotRecordSuccess() { + when(successMonitor.getRateAndReset()).thenReturn((double) mbpsToBps(100.0)); + when(failureMonitor.getRateAndReset()).thenReturn((double) mbpsToBps(10.0)); + regulator.regulate(); + EvictingQueue queue = getSuccessRateQueue(); + assertTrue(queue.isEmpty(), "Queue should be empty as failure occurred"); + } + + @SuppressWarnings("unchecked") + private EvictingQueue getSuccessRateQueue() { + Object field = getField(regulator, "successRateQueue"); + if (field instanceof EvictingQueue) { + return (EvictingQueue) field; + } + throw new IllegalStateException("Field 'successRateQueue' is not of expected type EvictingQueue"); + } + + /** + * Retrieves a static field value via reflection. + */ + private Object getStaticField(String fieldName) { + try { + Field field = TrafficRegulator.class.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(null); + } catch (Exception e) { + throw new RuntimeException("Unable to access " + fieldName + " field", e); + } + } + + private Object getField(Object instance, String fieldName) { + try { + Field field = TrafficRegulator.class.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(instance); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private int getWindowSize() { + return (int) getStaticField("WINDOW_SIZE"); + } + + private long getMinRateFromRegulator() { + return (long) getStaticField("MIN"); + } + + private long getMaxRateFromRegulator() { + return (long) getStaticField("MAX"); + } + + private double getFastIncrementRatio() { + return (double) getStaticField("FAST_INCREMENT_RATIO"); + } + + private double getSlowIncrementRatio() { + return (double) getStaticField("SLOW_INCREMENT_RATIO"); + } + + private void setRegulatorDecreaseEnv(double successRate, long limiterRate) { + setRegulatorEnv(successRate, FAILURE_RATE, limiterRate); + } + + private void setRegulatorIncreaseEnv(double successRate, long limiterRate) { + setRegulatorEnv(successRate, 0, limiterRate); + } + + private void setRegulatorEnv(double successRate, double failureRate, long limiterRate) { + when(successMonitor.getRateAndReset()).thenReturn(successRate); + when(failureMonitor.getRateAndReset()).thenReturn(failureRate); + when(rateLimiter.currentRate()).thenReturn(limiterRate); + } + + private void checkRegulate(long expectedNewRate) { + verify(rateLimiter).update(expectedNewRate); + verify(volumeLimiter).update(expectedNewRate * getWindowSize()); + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/streams/StreamManagerTest.java b/s3stream/src/test/java/com/automq/stream/s3/streams/StreamManagerTest.java index 8d4f8c7ebf..9699e4cee0 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/streams/StreamManagerTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/streams/StreamManagerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.streams; diff --git a/s3stream/src/test/java/com/automq/stream/s3/utils/AsyncRateLimiterTest.java b/s3stream/src/test/java/com/automq/stream/s3/utils/AsyncRateLimiterTest.java index d17ff60db8..8198f9f250 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/utils/AsyncRateLimiterTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/utils/AsyncRateLimiterTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.utils; diff --git a/s3stream/src/test/java/com/automq/stream/s3/utils/IdURITest.java b/s3stream/src/test/java/com/automq/stream/s3/utils/IdURITest.java index c96eb0ab75..79d71f9751 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/utils/IdURITest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/utils/IdURITest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.utils; diff --git a/s3stream/src/test/java/com/automq/stream/s3/utils/SecretUtilsTest.java b/s3stream/src/test/java/com/automq/stream/s3/utils/SecretUtilsTest.java index 16ab1c48d2..ba08f4fb60 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/utils/SecretUtilsTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/utils/SecretUtilsTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.s3.utils; diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockImplTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockImplTest.java deleted file mode 100644 index 09bf856b38..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockImplTest.java +++ /dev/null @@ -1,151 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.common.Record; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; - -import java.util.concurrent.CompletableFuture; - -import io.netty.buffer.ByteBuf; - -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; - -@Tag("S3Unit") -class BlockImplTest { - - static final int BODY_SIZE = 42; - static final int RECORD_SIZE = RECORD_HEADER_SIZE + BODY_SIZE; - static final Block.RecordSupplier RECORD_SUPPLIER = (offset, header) -> { - header.writerIndex(header.capacity()); - ByteBuf body = ByteBufAlloc.byteBuffer(BODY_SIZE); - body.writerIndex(body.capacity()); - return new Record(header, body); - }; - - private Block block; - - @AfterEach - void tearDown() { - if (block != null) { - block.release(); - } - } - - @Test - public void testAddRecord() { - block = new BlockImpl(0, Long.MAX_VALUE, Long.MAX_VALUE); - long offset; - offset = addRecord(block); - assertEquals(0, offset); - offset = addRecord(block); - assertEquals(RECORD_SIZE, offset); - assertNonEmptyBlock(block, 2); - } - - @Test - public void testExceedMaxSize() { - block = new BlockImpl(0, 1, Long.MAX_VALUE); - long offset; - offset = addRecord(block); - assertEquals(-1, offset, "Should return -1 when exceed max size"); - assertEmptyBlock(block); - } - - @Test - public void testExceedSoftLimit() { - block = new BlockImpl(0, Long.MAX_VALUE, RECORD_SIZE); - long offset; - offset = addRecord(block); - assertEquals(0, offset); - offset = addRecord(block); - assertEquals(-1, offset, "Should return -1 when exceed soft limit"); - assertNonEmptyBlock(block, 1); - } - - @Test - public void testOnlyOneRecordExceedSoftLimit() { - block = new BlockImpl(0, Long.MAX_VALUE, 1); - long offset; - offset = addRecord(block); - assertEquals(0, offset, "Should not fail when there is no record before, even exceed soft limit"); - assertNonEmptyBlock(block, 1); - } - - @Test - public void testFutures() { - block = new BlockImpl(0, Long.MAX_VALUE, Long.MAX_VALUE); - CompletableFuture future1 = new CompletableFuture<>(); - CompletableFuture future2 = new CompletableFuture<>(); - block.addRecord(RECORD_SIZE, RECORD_SUPPLIER, future1); - block.addRecord(RECORD_SIZE, RECORD_SUPPLIER, future2); - assertEquals(2, block.futures().size()); - assertTrue(block.futures().contains(future1)); - assertTrue(block.futures().contains(future2)); - } - - @Test - public void testCallDataTwice() { - block = new BlockImpl(0, Long.MAX_VALUE, Long.MAX_VALUE); - addRecord(block); - ByteBuf data1 = block.data(); - ByteBuf data2 = block.data(); - assertSame(data1, data2, "Should return the same data"); - } - - @Test - public void testRelease() { - block = new BlockImpl(0, Long.MAX_VALUE, Long.MAX_VALUE); - - ByteBuf body = ByteBufAlloc.byteBuffer(BODY_SIZE); - body.writerIndex(body.capacity()); - - Block.RecordSupplier recordSupplier = (offset, header) -> { - header.writerIndex(header.capacity()); - return new Record(header, body); - }; - block.addRecord(RECORD_SIZE, recordSupplier, new CompletableFuture<>()); - - ByteBuf data = block.data(); - block.release(); - - assertEquals(0, data.refCnt(), "Should release data"); - assertEquals(0, body.refCnt(), "Should release body"); - - // avoid double release - block = null; - } - - private static long addRecord(Block block) { - return block.addRecord(RECORD_SIZE, RECORD_SUPPLIER, new CompletableFuture<>()); - } - - private static void assertEmptyBlock(Block block) { - assertEquals(0, block.size()); - assertTrue(block.futures().isEmpty()); - assertEquals(0, block.data().readableBytes()); - } - - private static void assertNonEmptyBlock(Block block, int recordCount) { - assertEquals((long) recordCount * RECORD_SIZE, block.size()); - assertEquals(recordCount, block.futures().size()); - assertEquals(recordCount * RECORD_SIZE, block.data().readableBytes()); - } -} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALHeaderTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALHeaderTest.java deleted file mode 100644 index 007475852a..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALHeaderTest.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.wal.common.ShutdownType; -import com.automq.stream.s3.wal.exception.UnmarshalException; - -import org.junit.jupiter.api.Test; - -import static org.junit.jupiter.api.Assertions.assertEquals; - -public class BlockWALHeaderTest { - - @Test - public void test() throws UnmarshalException { - BlockWALHeader header = new BlockWALHeader(128 * 1024, 100); - header.updateTrimOffset(10); - header.setLastWriteTimestamp(11); - header.setShutdownType(ShutdownType.GRACEFULLY); - header.setNodeId(233); - header.setEpoch(234); - - BlockWALHeader unmarshal = BlockWALHeader.unmarshal(header.marshal().duplicate()); - assertEquals(header.getCapacity(), unmarshal.getCapacity()); - assertEquals(header.getTrimOffset(), unmarshal.getTrimOffset()); - assertEquals(header.getLastWriteTimestamp(), unmarshal.getLastWriteTimestamp()); - assertEquals(header.getSlidingWindowMaxLength(), unmarshal.getSlidingWindowMaxLength()); - assertEquals(header.getShutdownType(), unmarshal.getShutdownType()); - assertEquals(header.getNodeId(), unmarshal.getNodeId()); - assertEquals(header.getEpoch(), unmarshal.getEpoch()); - } - -} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java deleted file mode 100644 index b1fe87bee9..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/block/BlockWALServiceTest.java +++ /dev/null @@ -1,1386 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.block; - -import com.automq.stream.s3.ByteBufAlloc; -import com.automq.stream.s3.TestUtils; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.RecoverResult; -import com.automq.stream.s3.wal.WriteAheadLog; -import com.automq.stream.s3.wal.benchmark.WriteBench; -import com.automq.stream.s3.wal.common.RecordHeader; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.WALCapacityMismatchException; -import com.automq.stream.s3.wal.exception.WALNotInitializedException; -import com.automq.stream.s3.wal.impl.block.BlockWALService.RecoverIterator; -import com.automq.stream.s3.wal.util.WALBlockDeviceChannel; -import com.automq.stream.s3.wal.util.WALChannel; -import com.automq.stream.s3.wal.util.WALUtil; - -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.condition.EnabledOnOs; -import org.junit.jupiter.api.condition.OS; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.CsvSource; -import org.junit.jupiter.params.provider.MethodSource; -import org.junit.jupiter.params.provider.ValueSource; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Iterator; -import java.util.LinkedList; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; -import java.util.stream.Collectors; -import java.util.stream.Stream; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; -import io.netty.buffer.Unpooled; - -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_MAGIC_CODE; -import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; -import static com.automq.stream.s3.wal.impl.block.BlockWALService.WAL_HEADER_TOTAL_CAPACITY; -import static com.automq.stream.s3.wal.util.WALChannelTest.TEST_BLOCK_DEVICE_KEY; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -@Tag("S3Unit") -class BlockWALServiceTest { - - static final String TEST_BLOCK_DEVICE = System.getenv(TEST_BLOCK_DEVICE_KEY); - - private static void testSingleThreadAppendBasic0(boolean mergeWrite, - boolean directIO) throws IOException, OverCapacityException { - final int recordSize = 4096 + 1; - final int recordCount = 100; - final long blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount + WAL_HEADER_TOTAL_CAPACITY; - - String path = TestUtils.tempFilePath(); - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, blockDeviceCapacity); - } - - BlockWALService.BlockWALServiceBuilder builder = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .slidingWindowInitialSize(0) - .slidingWindowScaleUnit(4096); - if (!mergeWrite) { - builder.blockSoftLimit(0); - } - final WriteAheadLog wal = builder.build().start(); - recoverAndReset(wal); - - AtomicLong maxFlushedOffset = new AtomicLong(-1); - AtomicLong maxRecordOffset = new AtomicLong(-1); - try { - for (int i = 0; i < recordCount; i++) { - ByteBuf data = TestUtils.random(recordSize); - - final AppendResult appendResult = wal.append(data.retainedDuplicate()); - - if (!mergeWrite) { - assertEquals(0, appendResult.recordOffset() % WALUtil.BLOCK_SIZE); - } - appendResult.future().whenComplete((callbackResult, throwable) -> { - assertNull(throwable); - maxFlushedOffset.accumulateAndGet(callbackResult.flushedOffset(), Math::max); - maxRecordOffset.accumulateAndGet(appendResult.recordOffset(), Math::max); - if (!mergeWrite) { - assertEquals(0, callbackResult.flushedOffset() % WALUtil.alignLargeByBlockSize(recordSize)); - } else { - assertEquals(0, callbackResult.flushedOffset() % WALUtil.BLOCK_SIZE); - } - }).whenComplete((callbackResult, throwable) -> { - if (null != throwable) { - throwable.printStackTrace(); - Assertions.fail(); - } - }); - } - } finally { - wal.shutdownGracefully(); - } - assertTrue(maxFlushedOffset.get() > maxRecordOffset.get(), - "maxFlushedOffset should be greater than maxRecordOffset. maxFlushedOffset: " + maxFlushedOffset.get() + ", maxRecordOffset: " + maxRecordOffset.get()); - } - - private static void testSingleThreadAppendWhenOverCapacity0(boolean mergeWrite, - boolean directIO) throws IOException { - final int recordSize = 4096 + 1; - final int recordCount = 100; - long blockDeviceCapacity; - if (!mergeWrite) { - blockDeviceCapacity = recordCount / 3 * WALUtil.alignLargeByBlockSize(recordSize) + WAL_HEADER_TOTAL_CAPACITY; - } else { - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize * recordCount / 3) + WAL_HEADER_TOTAL_CAPACITY; - } - - String path = TestUtils.tempFilePath(); - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(blockDeviceCapacity); - resetBlockDevice(path, blockDeviceCapacity); - } - - BlockWALService.BlockWALServiceBuilder builder = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .slidingWindowInitialSize(0) - .slidingWindowScaleUnit(4096); - if (!mergeWrite) { - builder.blockSoftLimit(0); - } - final WriteAheadLog wal = builder.build().start(); - recoverAndReset(wal); - - AtomicLong maxFlushedOffset = new AtomicLong(-1); - AtomicLong maxRecordOffset = new AtomicLong(-1); - try { - WriteBench.TrimOffset trimOffset = new WriteBench.TrimOffset(); - for (int i = 0; i < recordCount; i++) { - ByteBuf data = TestUtils.random(recordSize); - AppendResult appendResult; - - while (true) { - try { - appendResult = wal.append(data.retainedDuplicate()); - } catch (OverCapacityException e) { - Thread.yield(); - wal.trim(trimOffset.get()).join(); - continue; - } - break; - } - - final long recordOffset = appendResult.recordOffset(); - if (!mergeWrite) { - assertEquals(0, recordOffset % WALUtil.BLOCK_SIZE); - } - trimOffset.appended(recordOffset); - appendResult.future().whenComplete((callbackResult, throwable) -> { - assertNull(throwable); - maxFlushedOffset.accumulateAndGet(callbackResult.flushedOffset(), Math::max); - maxRecordOffset.accumulateAndGet(recordOffset, Math::max); - if (!mergeWrite) { - assertEquals(0, callbackResult.flushedOffset() % WALUtil.alignLargeByBlockSize(recordSize)); - } else { - assertEquals(0, callbackResult.flushedOffset() % WALUtil.BLOCK_SIZE); - } - - trimOffset.flushed(callbackResult.flushedOffset()); - }).whenComplete((callbackResult, throwable) -> { - if (null != throwable) { - throwable.printStackTrace(); - Assertions.fail(); - } - }); - } - } finally { - wal.shutdownGracefully(); - } - assertTrue(maxFlushedOffset.get() > maxRecordOffset.get(), - "maxFlushedOffset should be greater than maxRecordOffset. maxFlushedOffset: " + maxFlushedOffset.get() + ", maxRecordOffset: " + maxRecordOffset.get()); - } - - private static void testMultiThreadAppend0(boolean mergeWrite, - boolean directIO) throws IOException, InterruptedException { - final int recordSize = 4096 + 1; - final int recordCount = 10; - final int threadCount = 8; - final long blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount * threadCount + WAL_HEADER_TOTAL_CAPACITY; - - String path = TestUtils.tempFilePath(); - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, blockDeviceCapacity); - } - - BlockWALService.BlockWALServiceBuilder builder = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO); - if (!mergeWrite) { - builder.blockSoftLimit(0); - } - final WriteAheadLog wal = builder.build().start(); - recoverAndReset(wal); - - ExecutorService executorService = Executors.newFixedThreadPool(threadCount); - AtomicLong maxFlushedOffset = new AtomicLong(-1); - AtomicLong maxRecordOffset = new AtomicLong(-1); - try { - for (int t = 0; t < threadCount; t++) { - executorService.submit(() -> Assertions.assertDoesNotThrow(() -> { - for (int i = 0; i < recordCount; i++) { - ByteBuf data = TestUtils.random(recordSize); - - final AppendResult appendResult = wal.append(data.retainedDuplicate()); - - appendResult.future().whenComplete((callbackResult, throwable) -> { - assertNull(throwable); - if (!mergeWrite) { - assertEquals(0, appendResult.recordOffset() % WALUtil.BLOCK_SIZE); - } - maxFlushedOffset.accumulateAndGet(callbackResult.flushedOffset(), Math::max); - maxRecordOffset.accumulateAndGet(appendResult.recordOffset(), Math::max); - if (!mergeWrite) { - assertEquals(0, callbackResult.flushedOffset() % WALUtil.alignLargeByBlockSize(recordSize)); - } else { - assertEquals(0, callbackResult.flushedOffset() % WALUtil.BLOCK_SIZE); - } - }).whenComplete((callbackResult, throwable) -> { - if (null != throwable) { - throwable.printStackTrace(); - Assertions.fail(); - } - }); - } - })); - } - } finally { - executorService.shutdown(); - assertTrue(executorService.awaitTermination(15, TimeUnit.SECONDS)); - wal.shutdownGracefully(); - } - assertTrue(maxFlushedOffset.get() > maxRecordOffset.get(), - "maxFlushedOffset should be greater than maxRecordOffset. maxFlushedOffset: " + maxFlushedOffset.get() + ", maxRecordOffset: " + maxRecordOffset.get()); - } - - private static void testRecoverAfterMergeWrite0(boolean shutdown, boolean overCapacity, - boolean directIO) throws IOException { - final int recordSize = 1024 + 1; - final int recordCount = 100; - long blockDeviceCapacity; - if (overCapacity) { - blockDeviceCapacity = recordSize * recordCount + WAL_HEADER_TOTAL_CAPACITY; - } else { - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount + WAL_HEADER_TOTAL_CAPACITY; - } - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(blockDeviceCapacity); - resetBlockDevice(path, blockDeviceCapacity); - } - - // Append records - final WriteAheadLog previousWAL = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - recoverAndReset(previousWAL); - List appended = appendAsync(previousWAL, recordSize, recordCount); - if (shutdown) { - previousWAL.shutdownGracefully(); - } - - // Recover records - final WriteAheadLog wal = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - try { - Iterator recover = recover(wal); - List recovered = new ArrayList<>(recordCount); - while (recover.hasNext()) { - RecoverResult next = recover.next(); - next.record().release(); - recovered.add(next.recordOffset()); - } - assertEquals(appended, recovered); - wal.reset().join(); - } finally { - wal.shutdownGracefully(); - } - } - - private static List appendAsync(WriteAheadLog wal, int recordSize, int recordCount) { - List appended = new ArrayList<>(recordCount); - List> appendFutures = new LinkedList<>(); - WriteBench.TrimOffset trimOffset = new WriteBench.TrimOffset(); - for (int i = 0; i < recordCount; i++) { - ByteBuf data = TestUtils.random(recordSize); - AppendResult appendResult; - try { - appendResult = wal.append(data.retainedDuplicate()); - } catch (OverCapacityException e) { - long offset = trimOffset.get(); - wal.trim(offset).join(); - appended = appended.stream() - .filter(recordOffset -> recordOffset > offset) - .collect(Collectors.toList()); - i--; - continue; - } - appended.add(appendResult.recordOffset()); - trimOffset.appended(appendResult.recordOffset()); - appendFutures.add(appendResult.future().whenComplete((callbackResult, throwable) -> { - assertNull(throwable); - assertEquals(0, callbackResult.flushedOffset() % WALUtil.BLOCK_SIZE); - trimOffset.flushed(callbackResult.flushedOffset()); - }).whenComplete((callbackResult, throwable) -> { - if (null != throwable) { - throwable.printStackTrace(); - Assertions.fail(); - } - }).thenApply(ignored -> null)); - } - CompletableFuture.allOf(appendFutures.toArray(new CompletableFuture[0])).join(); - return appended; - } - - public static Stream testRecoverFromDisasterData() { - return Stream.of( - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - -1L, - 50L, - Arrays.asList(0L, 2L, 4L), - Arrays.asList(0L, 2L, 4L), - WALUtil.BLOCK_SIZE - ).toArguments("base"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 0L, - 50L, - Arrays.asList(0L, 2L, 4L), - Arrays.asList(2L, 4L), - WALUtil.BLOCK_SIZE - ).toArguments("trimmed at zero"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 2L, - 50L, - Arrays.asList(0L, 2L, 4L, 6L), - Arrays.asList(4L, 6L), - WALUtil.BLOCK_SIZE - ).toArguments("trimmed"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 2L, - 50L, - Arrays.asList(0L, 2L, 4L, 6L, 8L, 10L, 12L, 14L, 16L, 18L, 20L), - Arrays.asList(4L, 6L, 8L, 10L, 12L, 14L, 16L, 18L, 20L), - WALUtil.BLOCK_SIZE - ).toArguments("WAL header flushed slow"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 2L, - 50L, - Arrays.asList(0L, 2L, 8L, 10L, 14L, 20L), - Arrays.asList(8L, 10L, 14L, 20L), - WALUtil.BLOCK_SIZE - ).toArguments("many invalid records"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 2L, - 50L, - Arrays.asList(14L, 8L, 10L, 20L, 0L, 2L), - Arrays.asList(8L, 10L, 14L, 20L), - WALUtil.BLOCK_SIZE - ).toArguments("write in random order"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 20230920L, - 50L, - Arrays.asList(20230900L, 20230910L, 20230916L, 20230920L, 20230930L, 20230940L, 20230950L, 20230960L, 20230970L, 20230980L), - Arrays.asList(20230930L, 20230940L, 20230950L, 20230960L, 20230970L), - WALUtil.BLOCK_SIZE - ).toArguments("big logic offset"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 180L, - 30L, - Arrays.asList(150L, 160L, 170L, 180L, 190L, 200L, 202L, 210L, 220L, 230L, 240L), - Arrays.asList(190L, 200L, 202L, 210L, 220L, 230L), - WALUtil.BLOCK_SIZE - ).toArguments("round robin"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE * 2 + 1, - 100L, - 192L, - 3L, - Arrays.asList(192L, 195L, /* no place for 198L, */ 200L, 203L, 206L, 209L, 212L, 215L), - Arrays.asList(195L, 200L, 203L, 206L, 209L, 212L, 215L), - WALUtil.BLOCK_SIZE - ).toArguments("round robin - no place for the last record"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 210L, - 50L, - Arrays.asList(111L, 113L, 115L, 117L, 119L, 120L, 130L, - 210L, 215L, 220L, 230L, 240L, 250L, 260L, 270L, 280L, 290L), - Arrays.asList(215L, 220L, 230L, 240L, 250L, 260L), - WALUtil.BLOCK_SIZE - ).toArguments("overwrite"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - -1L, - 1L, - Arrays.asList(0L, 2L, 5L, 7L), - List.of(0L, 2L), - WALUtil.BLOCK_SIZE - ).toArguments("small window - record size not aligned"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 10L, - 3L, - Arrays.asList(10L, 12L, 15L, 17L, 19L), - List.of(12L, 15L), - WALUtil.BLOCK_SIZE - ).toArguments("invalid record in window - record size not aligned"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE + 1, - 100L, - 10L, - 9L, - Arrays.asList(9L, 14L, 18L, 20L), - List.of(14L, 18L), - WALUtil.BLOCK_SIZE - ).toArguments("trim at an invalid record - record size not aligned"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE, - 100L, - -1L, - 1L, - Arrays.asList(0L, 1L, 3L, 4L, 5L), - List.of(0L, 1L), - WALUtil.BLOCK_SIZE - ).toArguments("small window - record size aligned"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE, - 100L, - 10L, - 3L, - Arrays.asList(10L, 11L, 13L, 14L, 15L, 16L), - List.of(11L, 13L, 14L), - WALUtil.BLOCK_SIZE - ).toArguments("invalid record in window - record size aligned"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE, - 100L, - 10L, - 5L, - Arrays.asList(9L, 11L, 13L, 15L, 16L, 17L), - List.of(11L, 13L, 15L, 16L), - WALUtil.BLOCK_SIZE - ).toArguments("trim at an invalid record - record size aligned"), - new RecoverFromDisasterParam( - WALUtil.BLOCK_SIZE, - 100L, - 10L, - 0L, - Arrays.asList(10L, 11L, 12L, 14L), - List.of(11L, 12L), - WALUtil.BLOCK_SIZE - ).toArguments("zero window"), - new RecoverFromDisasterParam( - 42, - 8192L, - -1L, - 8192L, - Arrays.asList(0L, 42L, 84L), - Arrays.asList(0L, 42L, 84L), - 1 - ).toArguments("merge write - base"), - new RecoverFromDisasterParam( - 42, - 8192L, - 42L, - 8192L, - Arrays.asList(0L, 42L, 84L, 126L), - Arrays.asList(84L, 126L), - 1 - ).toArguments("merge write - trimmed"), - new RecoverFromDisasterParam( - 42, - 8192L, - 42L, - 8192L, - Arrays.asList(0L, 42L, 42 * 2L, 42 * 4L, 4096L, 4096L + 42L, 4096L + 42 * 3L), - Arrays.asList(42 * 2L, 4096L, 4096L + 42L), - 1 - ).toArguments("merge write - some invalid records"), - new RecoverFromDisasterParam( - 42, - 8192L, - 42L, - 8192L, - Arrays.asList(42L, 42 * 4L, 42 * 2L, 4096L + 42 * 3L, 0L, 4096L, 4096L + 42L), - Arrays.asList(42 * 2L, 4096L, 4096L + 42L), - 1 - ).toArguments("merge write - random order"), - new RecoverFromDisasterParam( - 1000, - 8192L, - 2000L, - 8192L, - Arrays.asList(0L, 1000L, 2000L, 3000L, 4000L, 5000L, 7000L), - Arrays.asList(3000L, 4000L, 5000L), - 1 - ).toArguments("merge write - record in the middle"), - new RecoverFromDisasterParam( - 42, - 8192L, - 8192L + 4096L + 42L, - 8192L, - Arrays.asList(8192L + 4096L, 8192L + 4096L + 42L, 8192L + 4096L + 42 * 2L, 8192L + 4096L + 42 * 4L, 16384L, 16384L + 42L, 16384L + 42 * 3L), - Arrays.asList(8192L + 4096L + 42 * 2L, 16384L, 16384L + 42L), - 1 - ).toArguments("merge write - round robin"), - new RecoverFromDisasterParam( - 1000, - 8192L, - 12000L, - 8192L, - Arrays.asList(1000L, 2000L, 3000L, 4000L, 5000L, 6000L, 7000L, - 9000L, 10000L, 11000L, 12000L, 13000L, 14000L, 15000L), - Arrays.asList(13000L, 14000L, 15000L), - 1 - ).toArguments("merge write - overwrite"), - new RecoverFromDisasterParam( - 42, - 4096L * 20, - -1L, - 4096L, - Arrays.asList(0L, 42L, 42 * 3L, 4096L, 4096L + 42L, 4096L + 42 * 3L, 12288L, 12288L + 42L, 12288L + 42 * 3L, 16384L), - Arrays.asList(0L, 42L, 4096L, 4096L + 42L), - 1 - ).toArguments("merge write - small window"), - new RecoverFromDisasterParam( - 42, - 4096L * 20, - 4096L * 2, - 4096L * 4, - Arrays.asList(4096L * 2, 4096L * 2 + 42L, 4096L * 2 + 42 * 3L, - 4096L * 4, 4096L * 4 + 42L, 4096L * 4 + 42 * 3L, - 4096L * 5, 4096L * 5 + 42L, 4096L * 5 + 42 * 3L, - 4096L * 6, 4096L * 6 + 42L, 4096L * 6 + 42 * 3L, - 4096L * 7, 4096L * 7 + 42L, 4096L * 7 + 42 * 3L, - 4096L * 8), - Arrays.asList(4096L * 2 + 42L, 4096L * 4, 4096L * 4 + 42L, - 4096L * 5, 4096L * 5 + 42L, 4096L * 6, 4096L * 6 + 42L), - 1 - ).toArguments("merge write - invalid record in window"), - new RecoverFromDisasterParam( - 42, - 4096L * 20, - 4096L * 2 + 42 * 2L, - 4096L * 2, - Arrays.asList(4096L * 2, 4096L * 2 + 42L, 4096L * 2 + 42 * 3L, - 4096L * 3, 4096L * 3 + 42L, 4096L * 3 + 42 * 3L, - 4096L * 5, 4096L * 5 + 42L, 4096L * 5 + 42 * 3L, - 4096L * 6, 4096L * 6 + 42L, 4096L * 6 + 42 * 3L, - 4096L * 7), - Arrays.asList(4096L * 3, 4096L * 3 + 42L, 4096L * 5, 4096L * 5 + 42L), - 1 - ).toArguments("merge write - trim at an invalid record"), - new RecoverFromDisasterParam( - 42, - 4096L * 20, - 4096L * 2, - 0L, - Arrays.asList(4096L * 2, 4096L * 2 + 42L, 4096L * 2 + 42 * 3L, - 4096L * 3, 4096L * 3 + 42L, 4096L * 3 + 42 * 3L, - 4096L * 5, 4096L * 5 + 42L, 4096L * 5 + 42 * 3L, - 4096L * 6), - Arrays.asList(4096L * 2 + 42L, 4096L * 3, 4096L * 3 + 42L), - 1 - ).toArguments("merge write - zero window") - ); - } - - /** - * Call {@link WriteAheadLog#recover()} and set to strict mode. - */ - private static Iterator recover(WriteAheadLog wal) { - Iterator iterator = wal.recover(); - assertNotNull(iterator); - if (iterator instanceof RecoverIterator) { - ((RecoverIterator) iterator).strictMode(); - } - return iterator; - } - - /** - * Call {@link WriteAheadLog#recover()} {@link WriteAheadLog#reset()} and drop all records. - */ - private static void recoverAndReset(WriteAheadLog wal) { - for (Iterator it = recover(wal); it.hasNext(); ) { - it.next().record().release(); - } - wal.reset().join(); - } - - /** - * Write "0"s to the block device to reset it. - */ - private static void resetBlockDevice(String path, long capacity) throws IOException { - WALChannel channel = WALChannel.builder(path) - .capacity(capacity) - .direct(true) - .build(); - channel.open(); - ByteBuf buf = Unpooled.buffer((int) capacity); - buf.writeZero((int) capacity); - channel.write(buf, 0); - channel.close(); - } - - @ParameterizedTest(name = "Test {index}: mergeWrite={0}") - @ValueSource(booleans = {false, true}) - public void testSingleThreadAppendBasic(boolean mergeWrite) throws IOException, OverCapacityException { - testSingleThreadAppendBasic0(mergeWrite, false); - } - - @ParameterizedTest(name = "Test {index}: mergeWrite={0}") - @ValueSource(booleans = {false, true}) - @EnabledOnOs(OS.LINUX) - public void testSingleThreadAppendBasicDirectIO(boolean mergeWrite) throws IOException, OverCapacityException { - testSingleThreadAppendBasic0(mergeWrite, true); - } - - @ParameterizedTest(name = "Test {index}: mergeWrite={0}") - @ValueSource(booleans = {false, true}) - public void testSingleThreadAppendWhenOverCapacity(boolean mergeWrite) throws IOException { - testSingleThreadAppendWhenOverCapacity0(mergeWrite, false); - } - - @ParameterizedTest(name = "Test {index}: mergeWrite={0}") - @ValueSource(booleans = {false, true}) - @EnabledOnOs(OS.LINUX) - public void testSingleThreadAppendWhenOverCapacityDirectIO(boolean mergeWrite) throws IOException { - testSingleThreadAppendWhenOverCapacity0(mergeWrite, true); - } - - @ParameterizedTest(name = "Test {index}: mergeWrite={0}") - @ValueSource(booleans = {false, true}) - public void testMultiThreadAppend(boolean mergeWrite) throws InterruptedException, IOException { - testMultiThreadAppend0(mergeWrite, false); - } - - @ParameterizedTest(name = "Test {index}: mergeWrite={0}") - @ValueSource(booleans = {false, true}) - @EnabledOnOs(OS.LINUX) - public void testMultiThreadAppendDirectIO(boolean mergeWrite) throws InterruptedException, IOException { - testMultiThreadAppend0(mergeWrite, true); - } - - private long append(WriteAheadLog wal, int recordSize) throws OverCapacityException { - final AppendResult appendResult = wal.append(TestUtils.random(recordSize)); - final long recordOffset = appendResult.recordOffset(); - assertEquals(0, recordOffset % WALUtil.BLOCK_SIZE); - appendResult.future().whenComplete((callbackResult, throwable) -> { - assertNull(throwable); - assertTrue(callbackResult.flushedOffset() > recordOffset, "flushedOffset: " + callbackResult.flushedOffset() + ", recordOffset: " + recordOffset); - assertEquals(0, callbackResult.flushedOffset() % WALUtil.BLOCK_SIZE); - }).join(); - return recordOffset; - } - - private List append(WriteAheadLog wal, int recordSize, int recordCount) { - List recordOffsets = new ArrayList<>(recordCount); - long offset = 0; - for (int i = 0; i < recordCount; i++) { - try { - offset = append(wal, recordSize); - recordOffsets.add(offset); - } catch (OverCapacityException e) { - wal.trim(offset).join(); - final long trimmedOffset = offset; - recordOffsets = recordOffsets.stream() - .filter(recordOffset -> recordOffset > trimmedOffset) - .collect(Collectors.toList()); - i--; - } - } - return recordOffsets; - } - - private List appendWithoutTrim(WriteAheadLog wal, int recordSize, - int recordCount) throws OverCapacityException { - List recordOffsets = new ArrayList<>(recordCount); - for (int i = 0; i < recordCount; i++) { - long offset = append(wal, recordSize); - recordOffsets.add(offset); - } - return recordOffsets; - } - - @ParameterizedTest(name = "Test {index}: shutdown={0}, overCapacity={1}, recordCount={2}") - @CsvSource({ - "true, false, 10", - "true, true, 9", - "true, true, 10", - "true, true, 11", - - "false, false, 10", - "false, true, 9", - "false, true, 10", - "false, true, 11", - }) - public void testSingleThreadRecover(boolean shutdown, boolean overCapacity, int recordCount) throws IOException { - testSingleThreadRecover0(shutdown, overCapacity, recordCount, false); - } - - @ParameterizedTest(name = "Test {index}: shutdown={0}, overCapacity={1}, recordCount={2}") - @CsvSource({ - "true, false, 10", - "true, true, 9", - "true, true, 10", - "true, true, 11", - - "false, false, 10", - "false, true, 9", - "false, true, 10", - "false, true, 11", - }) - @EnabledOnOs(OS.LINUX) - public void testSingleThreadRecoverDirectIO(boolean shutdown, boolean overCapacity, - int recordCount) throws IOException { - testSingleThreadRecover0(shutdown, overCapacity, recordCount, true); - } - - private void testSingleThreadRecover0(boolean shutdown, boolean overCapacity, int recordCount, - boolean directIO) throws IOException { - final int recordSize = 4096 + 1; - long blockDeviceCapacity; - if (overCapacity) { - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount / 3 + WAL_HEADER_TOTAL_CAPACITY; - } else { - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount + WAL_HEADER_TOTAL_CAPACITY; - } - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(blockDeviceCapacity); - resetBlockDevice(path, blockDeviceCapacity); - } - - // Append records - final WriteAheadLog previousWAL = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - recoverAndReset(previousWAL); - List appended = append(previousWAL, recordSize, recordCount); - if (shutdown) { - previousWAL.shutdownGracefully(); - } - - // Recover records - final WriteAheadLog wal = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - try { - Iterator recover = recover(wal); - List recovered = new ArrayList<>(recordCount); - while (recover.hasNext()) { - RecoverResult next = recover.next(); - next.record().release(); - recovered.add(next.recordOffset()); - } - assertEquals(appended, recovered); - wal.reset().join(); - } finally { - wal.shutdownGracefully(); - } - } - - @ParameterizedTest(name = "Test {index}: shutdown={0}, overCapacity={1}") - @CsvSource({ - "true, false", - "true, true", - "false, false", - "false, true", - }) - public void testRecoverAfterMergeWrite(boolean shutdown, boolean overCapacity) throws IOException { - testRecoverAfterMergeWrite0(shutdown, overCapacity, false); - } - - @ParameterizedTest(name = "Test {index}: shutdown={0}, overCapacity={1}") - @CsvSource({ - "true, false", - "true, true", - "false, false", - "false, true", - }) - @EnabledOnOs(OS.LINUX) - public void testRecoverAfterMergeWriteDirectIO(boolean shutdown, boolean overCapacity) throws IOException { - testRecoverAfterMergeWrite0(shutdown, overCapacity, true); - } - - @Test - public void testAppendAfterRecover() throws IOException, OverCapacityException { - testAppendAfterRecover0(false); - } - - @Test - @EnabledOnOs(OS.LINUX) - public void testAppendAfterRecoverDirectIO() throws IOException, OverCapacityException { - testAppendAfterRecover0(true); - } - - private void testAppendAfterRecover0(boolean directIO) throws IOException, OverCapacityException { - final int recordSize = 4096 + 1; - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, 1 << 20); - } - - final WriteAheadLog previousWAL = BlockWALService.builder(path, 1 << 20) - .direct(directIO) - .build() - .start(); - recoverAndReset(previousWAL); - // Append 2 records - long appended0 = append(previousWAL, recordSize); - long appended1 = append(previousWAL, recordSize); - - final WriteAheadLog wal = BlockWALService.builder(path, 1 << 20) - .direct(directIO) - .build() - .start(); - try { - // Recover records - Iterator recover = recover(wal); - List recovered = new ArrayList<>(); - while (recover.hasNext()) { - RecoverResult next = recover.next(); - next.record().release(); - recovered.add(next.recordOffset()); - } - assertEquals(Arrays.asList(appended0, appended1), recovered); - - // Reset after recover - wal.reset().join(); - - // Append another 2 records - long appended2 = append(wal, recordSize); - long appended3 = append(wal, recordSize); - assertEquals(WALUtil.alignLargeByBlockSize(recordSize) + appended2, appended3); - } finally { - wal.shutdownGracefully(); - } - } - - private ByteBuf recordHeader(ByteBuf body, long offset) { - return new RecordHeader() - .setMagicCode(RECORD_HEADER_MAGIC_CODE) - .setRecordBodyLength(body.readableBytes()) - .setRecordBodyOffset(offset + RECORD_HEADER_SIZE) - .setRecordBodyCRC(WALUtil.crc32(body)) - .marshal(ByteBufAlloc.byteBuffer(RECORD_HEADER_SIZE), true); - } - - private void write(WALChannel walChannel, long logicOffset, int recordSize) throws IOException { - ByteBuf recordBody = TestUtils.random(recordSize - RECORD_HEADER_SIZE); - ByteBuf recordHeader = recordHeader(recordBody, logicOffset); - - CompositeByteBuf record = ByteBufAlloc.compositeByteBuffer(); - record.addComponents(true, recordHeader, recordBody); - - long position = WALUtil.recordOffsetToPosition(logicOffset, walChannel.capacity(), WAL_HEADER_TOTAL_CAPACITY); - writeAndFlush(walChannel, record, position); - } - - private void writeWALHeader(WALChannel walChannel, long trimOffset, long maxLength) throws IOException { - ByteBuf header = new BlockWALHeader(walChannel.capacity(), maxLength) - .updateTrimOffset(trimOffset) - .marshal(); - writeAndFlush(walChannel, header, 0); - } - - @ParameterizedTest(name = "Test {index} {0}") - @MethodSource("testRecoverFromDisasterData") - public void testRecoverFromDisaster( - String name, - int recordSize, - long capacity, - long trimOffset, - long maxLength, - List writeOffsets, - List recoveredOffsets - ) throws IOException { - testRecoverFromDisaster0(name, recordSize, capacity, trimOffset, maxLength, writeOffsets, recoveredOffsets, false); - } - - @ParameterizedTest(name = "Test {index} {0}") - @MethodSource("testRecoverFromDisasterData") - @EnabledOnOs({OS.LINUX}) - public void testRecoverFromDisasterDirectIO( - String name, - int recordSize, - long capacity, - long trimOffset, - long maxLength, - List writeOffsets, - List recoveredOffsets - ) throws IOException { - testRecoverFromDisaster0(name, recordSize, capacity, trimOffset, maxLength, writeOffsets, recoveredOffsets, true); - } - - private void testRecoverFromDisaster0( - String name, - int recordSize, - long capacity, - long trimOffset, - long maxLength, - List writeOffsets, - List recoveredOffsets, - boolean directIO - ) throws IOException { - String path = TestUtils.tempFilePath(); - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - capacity = WALUtil.alignLargeByBlockSize(capacity); - resetBlockDevice(path, capacity); - } - - WALChannel walChannel; - if (directIO) { - WALBlockDeviceChannel blockDeviceChannel = new WALBlockDeviceChannel(path, capacity); - blockDeviceChannel.unalignedWrite = true; - walChannel = blockDeviceChannel; - } else { - walChannel = WALChannel.builder(path) - .capacity(capacity) - .direct(false) - .build(); - } - - // Simulate disaster - walChannel.open(); - writeWALHeader(walChannel, trimOffset, maxLength); - for (long writeOffset : writeOffsets) { - write(walChannel, writeOffset, recordSize); - } - walChannel.close(); - - final WriteAheadLog wal = BlockWALService.builder(path, capacity) - .direct(directIO) - .build() - .start(); - try { - // Recover records - Iterator recover = recover(wal); - List recovered = new ArrayList<>(); - while (recover.hasNext()) { - RecoverResult next = recover.next(); - next.record().release(); - recovered.add(next.recordOffset()); - } - assertEquals(recoveredOffsets, recovered, name); - wal.reset().join(); - } finally { - wal.shutdownGracefully(); - } - } - - @Test - public void testRecoverAfterReset() throws IOException, OverCapacityException { - testRecoverAfterReset0(false); - } - - @Test - @EnabledOnOs({OS.LINUX}) - public void testRecoverAfterResetDirectIO() throws IOException, OverCapacityException { - testRecoverAfterReset0(true); - } - - private void testRecoverAfterReset0(boolean directIO) throws IOException, OverCapacityException { - final int recordSize = 4096 + 1; - final int recordCount = 10; - final long blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount * 2 + WAL_HEADER_TOTAL_CAPACITY; - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, blockDeviceCapacity); - } - - // 1. append and force shutdown - final WriteAheadLog wal1 = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - recoverAndReset(wal1); - List appended1 = appendWithoutTrim(wal1, recordSize, recordCount); - - // 2. recover and reset - final WriteAheadLog wal2 = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - Iterator recover = recover(wal2); - List recovered1 = new ArrayList<>(recordCount); - while (recover.hasNext()) { - RecoverResult next = recover.next(); - next.record().release(); - recovered1.add(next.recordOffset()); - } - assertEquals(appended1, recovered1); - wal2.reset().join(); - - // 3. append and force shutdown again - List appended2 = appendWithoutTrim(wal2, recordSize, recordCount); - - // 4. recover again - final WriteAheadLog wal3 = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - recover = recover(wal3); - List recovered2 = new ArrayList<>(recordCount); - while (recover.hasNext()) { - RecoverResult next = recover.next(); - next.record().release(); - recovered2.add(next.recordOffset()); - } - assertEquals(appended2, recovered2); - } - - @Test - public void testTrimInvalidOffset() throws IOException, OverCapacityException { - final WriteAheadLog wal = BlockWALService.builder(TestUtils.tempFilePath(), 16384) - .build() - .start(); - recoverAndReset(wal); - try { - long appended = append(wal, 42); - Assertions.assertThrows(IllegalArgumentException.class, () -> wal.trim(appended + 4096 + 1).join()); - } finally { - wal.shutdownGracefully(); - } - } - - @Test - public void testWindowGreaterThanCapacity() throws IOException, OverCapacityException { - final WriteAheadLog wal = BlockWALService.builder(TestUtils.tempFilePath(), WALUtil.BLOCK_SIZE * 3L) - .slidingWindowUpperLimit(WALUtil.BLOCK_SIZE * 4L) - .build() - .start(); - recoverAndReset(wal); - try { - append(wal, 42); - Assertions.assertThrows(OverCapacityException.class, () -> append(wal, 42)); - } finally { - wal.shutdownGracefully(); - } - } - - @Test - public void testRecoveryMode() throws IOException, OverCapacityException { - testRecoveryMode0(false); - } - - @Test - @EnabledOnOs({OS.LINUX}) - public void testRecoveryModeDirectIO() throws IOException, OverCapacityException { - testRecoveryMode0(true); - } - - private void testRecoveryMode0(boolean directIO) throws IOException, OverCapacityException { - final long capacity = 1 << 20; - final int nodeId = 10; - final long epoch = 100; - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, capacity); - } - - // simulate a crash - WriteAheadLog wal1 = BlockWALService.builder(path, capacity) - .nodeId(nodeId) - .epoch(epoch) - .direct(directIO) - .build() - .start(); - recoverAndReset(wal1); - wal1.append(TestUtils.random(4097)).future().join(); - - // open in recovery mode - WriteAheadLog wal2 = BlockWALService.recoveryBuilder(path) - .direct(directIO) - .build() - .start(); - assertEquals(nodeId, wal2.metadata().nodeId()); - assertEquals(epoch, wal2.metadata().epoch()); - // we can recover and reset the WAL - recoverAndReset(wal2); - // but we can't append to or trim it - assertThrows(IllegalStateException.class, () -> wal2.append(TestUtils.random(4097)).future().join()); - assertThrows(IllegalStateException.class, () -> wal2.trim(0).join()); - } - - @Test - public void testCapacityMismatchFileSize() throws IOException { - testCapacityMismatchFileSize0(false); - } - - @Test - @EnabledOnOs({OS.LINUX}) - public void testCapacityMismatchFileSizeDirectIO() throws IOException { - testCapacityMismatchFileSize0(true); - } - - private void testCapacityMismatchFileSize0(boolean directIO) throws IOException { - final long capacity1 = 1 << 20; - final long capacity2 = 1 << 21; - final String path = TestUtils.tempFilePath(); - - // init a WAL with capacity1 - WriteAheadLog wal1 = BlockWALService.builder(path, capacity1) - .direct(directIO) - .build() - .start(); - recoverAndReset(wal1); - wal1.shutdownGracefully(); - - // try to open it with capacity2 - WriteAheadLog wal2 = BlockWALService.builder(path, capacity2) - .direct(directIO) - .build(); - assertThrows(WALCapacityMismatchException.class, wal2::start); - } - - @Test - public void testCapacityMismatchInHeader() throws IOException { - testCapacityMismatchInHeader0(false); - } - - @Test - @EnabledOnOs({OS.LINUX}) - public void testCapacityMismatchInHeaderDirectIO() throws IOException { - testCapacityMismatchInHeader0(true); - } - - private void testCapacityMismatchInHeader0(boolean directIO) throws IOException { - final long capacity1 = 1 << 20; - final long capacity2 = 1 << 21; - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, capacity1); - } - - // init a WAL with capacity1 - WriteAheadLog wal1 = BlockWALService.builder(path, capacity1) - .direct(directIO) - .build() - .start(); - recoverAndReset(wal1); - wal1.shutdownGracefully(); - - // overwrite the capacity in the header with capacity2 - WALChannel walChannel = WALChannel.builder(path) - .capacity(capacity1) - .direct(directIO) - .build(); - walChannel.open(); - writeAndFlush(walChannel, new BlockWALHeader(capacity2, 42).marshal(), 0); - walChannel.close(); - - // try to open it with capacity1 - WriteAheadLog wal2 = BlockWALService.builder(path, capacity1) - .direct(directIO) - .build(); - assertThrows(WALCapacityMismatchException.class, wal2::start); - } - - @Test - public void testRecoveryModeWALFileNotExist() throws IOException { - testRecoveryModeWALFileNotExist0(false); - } - - @Test - @EnabledOnOs({OS.LINUX}) - public void testRecoveryModeWALFileNotExistDirectIO() throws IOException { - testRecoveryModeWALFileNotExist0(true); - } - - private void testRecoveryModeWALFileNotExist0(boolean directIO) throws IOException { - final String path = TestUtils.tempFilePath(); - - WriteAheadLog wal = BlockWALService.recoveryBuilder(path) - .direct(directIO) - .build(); - assertThrows(WALNotInitializedException.class, wal::start); - } - - @Test - public void testRecoveryModeNoHeader() throws IOException { - testRecoveryModeNoHeader0(false); - } - - @Test - @EnabledOnOs({OS.LINUX}) - public void testRecoveryModeNoHeaderDirectIO() throws IOException { - testRecoveryModeNoHeader0(true); - } - - private void testRecoveryModeNoHeader0(boolean directIO) throws IOException { - final long capacity = 1 << 20; - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - resetBlockDevice(path, capacity); - } - - // init a WAL - WriteAheadLog wal1 = BlockWALService.builder(path, capacity) - .direct(directIO) - .build() - .start(); - recoverAndReset(wal1); - wal1.shutdownGracefully(); - - // clear the WAL header - WALChannel walChannel = WALChannel.builder(path) - .capacity(capacity) - .direct(directIO) - .build(); - walChannel.open(); - writeAndFlush(walChannel, Unpooled.buffer(WAL_HEADER_TOTAL_CAPACITY).writeZero(WAL_HEADER_TOTAL_CAPACITY), 0); - walChannel.close(); - - // try to open it in recovery mode - WriteAheadLog wal2 = BlockWALService.recoveryBuilder(path) - .direct(directIO) - .build(); - assertThrows(WALNotInitializedException.class, wal2::start); - } - - @ParameterizedTest(name = "Test {index}: recordCount={0}") - @ValueSource(ints = {10, 20, 30, 40}) - public void testResetWithoutRecover(int recordCount) throws IOException { - testResetWithoutRecover0(recordCount, false); - } - - @ParameterizedTest(name = "Test {index}: recordCount={0}") - @ValueSource(ints = {10, 20, 30, 40}) - @EnabledOnOs({OS.LINUX}) - public void testResetWithoutRecoverDirectIO(int recordCount) throws IOException { - testResetWithoutRecover0(recordCount, true); - } - - private void testResetWithoutRecover0(int recordCount, boolean directIO) throws IOException { - final int recordSize = 4096 + 1; - long blockDeviceCapacity = WALUtil.alignLargeByBlockSize(recordSize) * recordCount / 3 + WAL_HEADER_TOTAL_CAPACITY; - String path = TestUtils.tempFilePath(); - - if (directIO && TEST_BLOCK_DEVICE != null) { - path = TEST_BLOCK_DEVICE; - blockDeviceCapacity = WALUtil.alignLargeByBlockSize(blockDeviceCapacity); - resetBlockDevice(path, blockDeviceCapacity); - } - - // Append records - final WriteAheadLog wal1 = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - recoverAndReset(wal1); - wal1.shutdownGracefully(); - - // Reset WAL without recover - final WriteAheadLog wal2 = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - try { - wal2.reset().join(); - } finally { - wal2.shutdownGracefully(); - } - - // Try to recover records - final WriteAheadLog wal3 = BlockWALService.builder(path, blockDeviceCapacity) - .direct(directIO) - .build() - .start(); - try { - Iterator recover = recover(wal3); - assertFalse(recover.hasNext()); - } finally { - wal3.shutdownGracefully(); - } - } - - private void writeAndFlush(WALChannel channel, ByteBuf src, long position) throws IOException { - channel.write(src, position); - channel.flush(); - } - - private static class RecoverFromDisasterParam { - int recordSize; - long capacity; - // WAL header - long trimOffset; - long maxLength; - // WAL records - List writeOffsets; - List recoveredOffsets; - - public RecoverFromDisasterParam( - int recordSize, - long capacity, - long trimOffset, - long maxLength, - List writeOffsets, - List recoveredOffsets, - int unit - ) { - this.recordSize = recordSize; - this.capacity = capacity * unit + WAL_HEADER_TOTAL_CAPACITY; - this.trimOffset = trimOffset * unit; - this.maxLength = maxLength * unit; - this.writeOffsets = writeOffsets.stream().map(offset -> offset * unit).collect(Collectors.toList()); - this.recoveredOffsets = recoveredOffsets.stream().map(offset -> offset * unit).collect(Collectors.toList()); - } - - public Arguments toArguments(String name) { - return Arguments.of(name, recordSize, capacity, trimOffset, maxLength, writeOffsets, recoveredOffsets); - } - } -} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/DefaultWriterTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/DefaultWriterTest.java new file mode 100644 index 0000000000..55a472ffdd --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/DefaultWriterTest.java @@ -0,0 +1,444 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.utils.Time; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; + +import java.util.Random; +import java.util.concurrent.ConcurrentSkipListMap; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +public class DefaultWriterTest { + private DefaultWriter writer; + private MockObjectStorage objectStorage; + private ConcurrentSkipListMap generatedByteBufMap; + private Random random; + + @BeforeEach + public void setUp() { + objectStorage = new MockObjectStorage(); + ObjectWALConfig config = ObjectWALConfig.builder() + .withMaxBytesInBatch(123) + .withNodeId(100) + .withEpoch(1000) + .withBatchInterval(Long.MAX_VALUE) + .build(); + writer = new DefaultWriter(Time.SYSTEM, objectStorage, config); + writer.start(); + generatedByteBufMap = new ConcurrentSkipListMap<>(); + random = new Random(); + } + + @AfterEach + public void tearDown() { + objectStorage.triggerAll(); + writer.close(); + objectStorage.close(); + } + + private ByteBuf generateByteBuf(int size) { + ByteBuf byteBuf = Unpooled.buffer(size); + byte[] bytes = new byte[size]; + random.nextBytes(bytes); + byteBuf.writeBytes(bytes); + return byteBuf; + } + + // TODO: fix the test +// @Test +// public void testOffset() throws OverCapacityException, WALFencedException { +// ByteBuf byteBuf1 = generateByteBuf(50); +// CompletableFuture future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf1.readableBytes(), offset -> byteBuf1.retainedSlice().asReadOnly(), future); +// assertEquals(50, recordAccumulator.nextOffset()); +// +// recordAccumulator.unsafeUpload(true); +// long flushedOffset = future.join().flushedOffset(); +// assertEquals(50, flushedOffset); +// assertEquals(50, recordAccumulator.flushedOffset()); +// +// List objectList = recordAccumulator.objectList(); +// assertEquals(1, objectList.size()); +// +// RecordAccumulator.WALObject object = objectList.get(0); +// assertEquals(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE + 50, object.length()); +// ByteBuf result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); +// ByteBuf headerBuf = result.readBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE); +// WALObjectHeader objectHeader = WALObjectHeader.unmarshal(headerBuf); +// headerBuf.release(); +// assertEquals(WALObjectHeader.DEFAULT_WAL_MAGIC_CODE, objectHeader.magicCode()); +// assertEquals(0, objectHeader.startOffset()); +// assertEquals(50, objectHeader.length()); +// // The last write timestamp is not set currently. +// assertEquals(0L, objectHeader.stickyRecordLength()); +// assertEquals(100, objectHeader.nodeId()); +// assertEquals(1000, objectHeader.epoch()); +// assertEquals(-1, objectHeader.trimOffset()); +// +// assertEquals(byteBuf1, result); +// byteBuf1.release(); +// +// // Test huge record. +// ByteBuf byteBuf2 = generateByteBuf(50); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf2.readableBytes(), offset -> byteBuf2.retainedSlice().asReadOnly(), future); +// assertEquals(100, recordAccumulator.nextOffset()); +// +// ByteBuf byteBuf3 = generateByteBuf(75); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf3.readableBytes(), offset -> byteBuf3.retainedSlice().asReadOnly(), future); +// assertEquals(175, recordAccumulator.nextOffset()); +// +// recordAccumulator.unsafeUpload(true); +// flushedOffset = future.join().flushedOffset(); +// assertEquals(175, flushedOffset); +// assertEquals(175, recordAccumulator.flushedOffset()); +// +// objectList = recordAccumulator.objectList(); +// assertEquals(2, objectList.size()); +// +// object = objectList.get(1); +// assertEquals(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE + 50 + 75, object.length()); +// result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); +// result.skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE); +// CompositeByteBuf compositeBuffer = Unpooled.compositeBuffer(); +// compositeBuffer.addComponents(true, byteBuf2); +// compositeBuffer.addComponents(true, byteBuf3); +// assertEquals(compositeBuffer, result); +// compositeBuffer.release(); +// +// // Test record part +// ByteBuf byteBuf4 = generateByteBuf(50); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf4.readableBytes(), offset -> byteBuf4.retainedSlice().asReadOnly(), future); +// assertEquals(225, recordAccumulator.nextOffset()); +// +// ByteBuf byteBuf5 = generateByteBuf(50); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf5.readableBytes(), offset -> byteBuf5.retainedSlice().asReadOnly(), future); +// assertEquals(275, recordAccumulator.nextOffset()); +// +// recordAccumulator.unsafeUpload(true); +// flushedOffset = future.join().flushedOffset(); +// assertEquals(275, flushedOffset); +// assertEquals(275, recordAccumulator.flushedOffset()); +// +// objectList = recordAccumulator.objectList(); +// assertEquals(4, objectList.size()); +// +// object = objectList.get(2); +// assertEquals(123, object.length()); +// result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); +// result.skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE); +// assertEquals(byteBuf4, result.readBytes(50)); +// +// object = objectList.get(3); +// compositeBuffer = Unpooled.compositeBuffer(); +// compositeBuffer.addComponents(true, result); +// result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); +// result.skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE); +// compositeBuffer.addComponents(true, result); +// assertEquals(compositeBuffer, byteBuf5); +// byteBuf4.release(); +// byteBuf5.release(); +// } +// +// @Test +// public void testStrictBatchLimit() throws OverCapacityException, WALFencedException { +// CompletableFuture future = new CompletableFuture<>(); +// recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); +// recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); +// recordAccumulator.append(50, offset -> generateByteBuf(50), future); +// assertEquals(150, recordAccumulator.nextOffset()); +// +// recordAccumulator.unsafeUpload(true); +// future.join(); +// +// assertEquals(2, recordAccumulator.objectList().size()); +// +// // Reset the RecordAccumulator with strict batch limit disabled. +// recordAccumulator.close(); +// ObjectWALConfig config = ObjectWALConfig.builder() +// .withMaxBytesInBatch(115) +// .withNodeId(100) +// .withEpoch(1000) +// .withBatchInterval(Long.MAX_VALUE) +// .withStrictBatchLimit(false) +// .build(); +// recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ReservationService.NOOP, config); +// recordAccumulator.start(); +// +// assertEquals(2, recordAccumulator.objectList().size()); +// +// future = new CompletableFuture<>(); +// recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); +// recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); +// recordAccumulator.append(50, offset -> generateByteBuf(50), future); +// assertEquals(300, recordAccumulator.nextOffset()); +// +// +// recordAccumulator.unsafeUpload(true); +// future.join(); +// +// assertEquals(3, recordAccumulator.objectList().size()); +// } +// +// @ParameterizedTest +// @ValueSource(booleans = {true, false}) +// public void testInMultiThread(boolean strictBathLimit) throws InterruptedException, WALFencedException { +// recordAccumulator.close(); +// +// ObjectWALConfig config = ObjectWALConfig.builder() +// .withMaxBytesInBatch(115) +// .withNodeId(100) +// .withEpoch(1000) +// .withBatchInterval(Long.MAX_VALUE) +// .withStrictBatchLimit(strictBathLimit) +// .build(); +// recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ReservationService.NOOP, config); +// recordAccumulator.start(); +// +// int threadCount = 10; +// CountDownLatch startBarrier = new CountDownLatch(threadCount); +// CountDownLatch stopCountDownLatch = new CountDownLatch(threadCount); +// List> futureList = new CopyOnWriteArrayList<>(); +// for (int i = 0; i < threadCount; i++) { +// new Thread(() -> { +// startBarrier.countDown(); +// try { +// startBarrier.await(); +// } catch (InterruptedException e) { +// throw new RuntimeException(e); +// } +// +// for (int j = 0; j < 100; j++) { +// ByteBuf byteBuf = generateByteBuf(40); +// try { +// CompletableFuture future = new CompletableFuture<>(); +// long offset = recordAccumulator.append(byteBuf.readableBytes(), o -> byteBuf.retainedSlice().asReadOnly(), future); +// futureList.add(future); +// generatedByteBufMap.put(offset, byteBuf); +// +// Thread.sleep(15); +// } catch (Exception e) { +// throw new RuntimeException(e); +// } +// } +// +// stopCountDownLatch.countDown(); +// }).start(); +// } +// +// stopCountDownLatch.await(); +// +// // Ensure all records are uploaded. +// try { +// recordAccumulator.unsafeUpload(true); +// } catch (Exception e) { +// fail(e); +// } +// +// for (CompletableFuture future : futureList) { +// future.join(); +// } +// +// assertEquals(100 * threadCount, generatedByteBufMap.size()); +// +// assertFalse(recordAccumulator.objectList().isEmpty()); +// +// CompositeByteBuf source = Unpooled.compositeBuffer(); +// for (ByteBuf buffer : generatedByteBufMap.values()) { +// source.addComponent(true, buffer); +// } +// +// CompositeByteBuf result = Unpooled.compositeBuffer(); +// for (RecordAccumulator.WALObject object : recordAccumulator.objectList()) { +// ByteBuf buf = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); +// buf.skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE); +// result.addComponent(true, buf); +// } +// +// assertEquals(source, result); +// source.release(); +// result.release(); +// } +// +// @Test +// public void testUploadPeriodically() throws OverCapacityException, WALFencedException { +// recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ReservationService.NOOP, ObjectWALConfig.builder().build()); +// recordAccumulator.start(); +// +// assertTrue(recordAccumulator.objectList().isEmpty()); +// +// ByteBuf byteBuf = generateByteBuf(25); +// CompletableFuture future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf.readableBytes(), o -> byteBuf.retainedSlice().asReadOnly(), future); +// +// await().atMost(Duration.ofSeconds(1)).until(future::isDone); +// assertEquals(1, recordAccumulator.objectList().size()); +// } +// +// @Test +// public void testShutdown() throws InterruptedException, OverCapacityException, WALFencedException { +// ScheduledExecutorService executorService = recordAccumulator.executorService(); +// executorService.shutdown(); +// executorService.awaitTermination(10, TimeUnit.SECONDS); +// +// ByteBuf byteBuf = generateByteBuf(25); +// CompletableFuture future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf.readableBytes(), o -> byteBuf.retainedSlice().asReadOnly(), future); +// +// await().during(Duration.ofSeconds(1)).atMost(Duration.ofSeconds(3)).until(() -> !future.isDone()); +// assertTrue(recordAccumulator.objectList().isEmpty()); +// +// // Flush all data to S3 when close. +// recordAccumulator.close(); +// assertTrue(future.isDone()); +// +// +// ObjectWALConfig config = ObjectWALConfig.builder() +// .withMaxBytesInBatch(115) +// .withNodeId(100) +// .withEpoch(1000) +// .withBatchInterval(Long.MAX_VALUE) +// .withStrictBatchLimit(true) +// .build(); +// recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ReservationService.NOOP, config); +// recordAccumulator.start(); +// assertEquals(1, recordAccumulator.objectList().size()); +// } +// +// @Test +// public void testTrim() throws OverCapacityException, WALFencedException { +// ByteBuf byteBuf1 = generateByteBuf(50); +// CompletableFuture future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf1.readableBytes(), offset -> byteBuf1.retainedSlice().asReadOnly(), future); +// +// ByteBuf byteBuf2 = generateByteBuf(50); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf2.readableBytes(), offset -> byteBuf2.retainedSlice().asReadOnly(), future); +// +// recordAccumulator.unsafeUpload(true); +// long flushedOffset = future.join().flushedOffset(); +// assertEquals(100, flushedOffset); +// assertEquals(100, recordAccumulator.flushedOffset()); +// assertEquals(2, recordAccumulator.objectList().size()); +// +// recordAccumulator.trim(50).join(); +// assertEquals(2, recordAccumulator.objectList().size()); +// +// recordAccumulator.trim(100).join(); +// assertEquals(0, recordAccumulator.objectList().size()); +// } +// +// @Test +// public void testReset() throws OverCapacityException, WALFencedException { +// ByteBuf byteBuf1 = generateByteBuf(50); +// CompletableFuture future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf1.readableBytes(), offset -> byteBuf1.retainedSlice().asReadOnly(), future); +// recordAccumulator.unsafeUpload(true); +// future.join(); +// +// ByteBuf byteBuf2 = generateByteBuf(50); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf2.readableBytes(), offset -> byteBuf2.retainedSlice().asReadOnly(), future); +// recordAccumulator.unsafeUpload(true); +// future.join(); +// +// // Close and restart with another node id. +// recordAccumulator.close(); +// recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ReservationService.NOOP, ObjectWALConfig.builder().withEpoch(System.currentTimeMillis()).build()); +// recordAccumulator.start(); +// assertEquals(0, recordAccumulator.objectList().size()); +// +// // Close and restart with the same node id and higher node epoch. +// recordAccumulator.close(); +// recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ReservationService.NOOP, ObjectWALConfig.builder().withNodeId(100).withEpoch(System.currentTimeMillis()).build()); +// recordAccumulator.start(); +// assertEquals(2, recordAccumulator.objectList().size()); +// +// ByteBuf byteBuf3 = generateByteBuf(50); +// future = new CompletableFuture<>(); +// recordAccumulator.append(byteBuf3.readableBytes(), offset -> byteBuf3.retainedSlice().asReadOnly(), future); +// recordAccumulator.unsafeUpload(true); +// future.join(); +// +// List objectList = recordAccumulator.objectList(); +// assertEquals(3, objectList.size()); +// assertEquals(byteBuf1, objectStorage.read(new ReadOptions().bucket((short) 0), objectList.get(0).path()).join().skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE)); +// assertEquals(byteBuf2, objectStorage.read(new ReadOptions().bucket((short) 0), objectList.get(1).path()).join().skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE)); +// assertEquals(byteBuf3, objectStorage.read(new ReadOptions().bucket((short) 0), objectList.get(2).path()).join().skipBytes(WALObjectHeader.DEFAULT_WAL_HEADER_SIZE)); +// +// recordAccumulator.reset().join(); +// assertEquals(0, recordAccumulator.objectList().size()); +// } +// +// @Test +// public void testSequentiallyComplete() throws WALFencedException, OverCapacityException, InterruptedException { +// objectStorage.markManualWrite(); +// ByteBuf byteBuf = generateByteBuf(1); +// +// CompletableFuture future0 = new CompletableFuture<>(); +// CompletableFuture future1 = new CompletableFuture<>(); +// CompletableFuture future2 = new CompletableFuture<>(); +// CompletableFuture future3 = new CompletableFuture<>(); +// +// recordAccumulator.append(byteBuf.readableBytes(), offset -> byteBuf.retainedSlice().asReadOnly(), future0); +// recordAccumulator.unsafeUpload(true); +// recordAccumulator.append(byteBuf.readableBytes(), offset -> byteBuf.retainedSlice().asReadOnly(), future1); +// recordAccumulator.append(byteBuf.readableBytes(), offset -> byteBuf.retainedSlice().asReadOnly(), future2); +// recordAccumulator.unsafeUpload(true); +// recordAccumulator.append(byteBuf.readableBytes(), offset -> byteBuf.retainedSlice().asReadOnly(), future3); +// recordAccumulator.unsafeUpload(true); +// +// // sleep to wait for potential async callback +// Thread.sleep(100); +// assertFalse(future0.isDone()); +// assertFalse(future1.isDone()); +// assertFalse(future2.isDone()); +// assertFalse(future3.isDone()); +// +// objectStorage.triggerWrite("1-3"); +// Thread.sleep(100); +// assertFalse(future0.isDone()); +// assertFalse(future1.isDone()); +// assertFalse(future2.isDone()); +// assertFalse(future3.isDone()); +// +// objectStorage.triggerWrite("0-1"); +// Thread.sleep(100); +// assertTrue(future0.isDone()); +// assertTrue(future1.isDone()); +// assertTrue(future2.isDone()); +// assertFalse(future3.isDone()); +// +// objectStorage.triggerWrite("3-4"); +// Thread.sleep(100); +// assertTrue(future0.isDone()); +// assertTrue(future1.isDone()); +// assertTrue(future2.isDone()); +// assertTrue(future3.isDone()); +// } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/MockObjectStorage.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/MockObjectStorage.java new file mode 100644 index 0000000000..aff09aab07 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/MockObjectStorage.java @@ -0,0 +1,68 @@ +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.operator.MemoryObjectStorage; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; + +import io.netty.buffer.ByteBuf; + +import static org.junit.jupiter.api.Assertions.assertNotNull; + +public class MockObjectStorage extends MemoryObjectStorage { + + private final Map pendingWrites = new ConcurrentHashMap<>(); + private boolean manualWrite = false; + + @Override + public CompletableFuture write(WriteOptions options, String objectPath, ByteBuf data) { + Supplier> pendingWrite = () -> super.write(options, objectPath, data); + if (!manualWrite) { + return pendingWrite.get(); + } + + CompletableFuture mockFuture = new CompletableFuture<>(); + String offset = objectPath.substring(objectPath.lastIndexOf('/') + 1); + pendingWrites.put(offset, new PendingWrite(mockFuture, pendingWrite)); + + return mockFuture; + } + + public void markManualWrite() { + manualWrite = true; + } + + public void triggerAll() { + pendingWrites.values().forEach(PendingWrite::trigger); + pendingWrites.clear(); + } + + public void triggerWrite(String objectPath) { + PendingWrite pendingWrite = pendingWrites.remove(objectPath); + assertNotNull(pendingWrite); + pendingWrite.trigger(); + } + + private static class PendingWrite { + public final CompletableFuture mockFuture; + public final Supplier> originalFuture; + + public PendingWrite(CompletableFuture mockFuture, + Supplier> originalFuture) { + this.mockFuture = mockFuture; + this.originalFuture = originalFuture; + } + + public void trigger() { + originalFuture.get().whenComplete((result, error) -> { + if (error != null) { + mockFuture.completeExceptionally(error); + } else { + mockFuture.complete(result); + } + }); + } + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectReservationServiceTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectReservationServiceTest.java new file mode 100644 index 0000000000..026ba17261 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectReservationServiceTest.java @@ -0,0 +1,97 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.s3.wal.impl.object; + +import com.automq.stream.s3.operator.MemoryObjectStorage; +import com.automq.stream.s3.operator.ObjectStorage; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class ObjectReservationServiceTest { + private ObjectReservationService reservationService; + private ObjectStorage objectStorage; + + @BeforeEach + public void setUp() { + objectStorage = new MemoryObjectStorage(); + reservationService = new ObjectReservationService("cluster", objectStorage, (short) 0); + } + + @Test + void verify() { + reservationService.acquire(1, 2, false).join(); + assertTrue(reservationService.verify(1, 2, false).join()); + assertFalse(reservationService.verify(1, 2, true).join()); + + assertFalse(reservationService.verify(1, 1, false).join()); + assertFalse(reservationService.verify(1, 3, false).join()); + assertFalse(reservationService.verify(2, 2, false).join()); + + reservationService.acquire(1, 2, true).join(); + assertFalse(reservationService.verify(1, 2, false).join()); + assertTrue(reservationService.verify(1, 2, true).join()); + } + + @Test + void acquire() { + ByteBuf target = Unpooled.buffer(Long.BYTES * 10); + target.writeLong(Long.MAX_VALUE); + target.writeInt(ObjectReservationService.S3_RESERVATION_OBJECT_MAGIC_CODE); + target.writeLong(1); + target.writeLong(2); + target.writeBoolean(false); + target.writeLong(Long.MAX_VALUE); + target.readerIndex(Long.BYTES); + target.writerIndex(Long.BYTES + ObjectReservationService.S3_RESERVATION_OBJECT_LENGTH); + reservationService.acquire(1, 2, false).join(); + assertTrue(reservationService.verify(1, target).join()); + + target = Unpooled.buffer(Long.BYTES * 10); + target.writeInt(ObjectReservationService.S3_RESERVATION_OBJECT_MAGIC_CODE); + target.writeLong(1); + target.writeLong(2); + target.writeBoolean(false); + reservationService.acquire(1, 2, true).join(); + assertFalse(reservationService.verify(1, target).join()); + + target = Unpooled.buffer(Long.BYTES * 10); + target.writeInt(ObjectReservationService.S3_RESERVATION_OBJECT_MAGIC_CODE); + target.writeLong(1); + target.writeLong(2); + target.writeBoolean(true); + reservationService.acquire(1, 2, true).join(); + assertTrue(reservationService.verify(1, target).join()); + + target = Unpooled.buffer(Long.BYTES * 10); + target.writeInt(ObjectReservationService.S3_RESERVATION_OBJECT_MAGIC_CODE); + target.writeLong(1); + target.writeLong(2); + target.writeBoolean(true); + reservationService.acquire(1, 2, false).join(); + assertFalse(reservationService.verify(1, target).join()); + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java index 0b74675b2d..0ef2f36b42 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALServiceTest.java @@ -1,135 +1,447 @@ package com.automq.stream.s3.wal.impl.object; -import com.automq.stream.s3.operator.MemoryObjectStorage; +import com.automq.stream.s3.ByteBufAlloc; +import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.trace.context.TraceContext; import com.automq.stream.s3.wal.AppendResult; import com.automq.stream.s3.wal.RecoverResult; -import com.automq.stream.s3.wal.common.RecordHeader; +import com.automq.stream.s3.wal.common.Record; import com.automq.stream.s3.wal.exception.OverCapacityException; +import com.automq.stream.s3.wal.impl.DefaultRecordOffset; +import com.automq.stream.s3.wal.util.WALUtil; +import com.automq.stream.utils.MockTime; import com.automq.stream.utils.Time; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import java.io.IOException; import java.util.ArrayList; -import java.util.Iterator; import java.util.List; import java.util.Random; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.stream.Stream; import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; import io.netty.buffer.Unpooled; +import static com.automq.stream.s3.wal.common.RecordHeader.RECORD_HEADER_SIZE; +import static com.automq.stream.s3.wal.impl.object.RecoverIterator.getContinuousFromTrimOffset; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; +@Timeout(120) public class ObjectWALServiceTest { - private ObjectStorage objectStorage; - private ObjectWALService wal; + private MockObjectStorage objectStorage; private Random random; + private Time time; @BeforeEach public void setUp() throws IOException { - objectStorage = new MemoryObjectStorage(); - ObjectWALConfig config = ObjectWALConfig.builder() - .withMaxBytesInBatch(110) - .withBatchInterval(Long.MAX_VALUE) - .withStrictBatchLimit(true) - .build(); - wal = new ObjectWALService(Time.SYSTEM, objectStorage, config); - wal.start(); + objectStorage = new MockObjectStorage(); + time = new MockTime(); random = new Random(); } @AfterEach public void tearDown() { - wal.shutdownGracefully(); + objectStorage.triggerAll(); objectStorage.close(); } - private ByteBuf generateByteBuf(int size) { - ByteBuf byteBuf = Unpooled.buffer(size); - byte[] bytes = new byte[size]; - random.nextBytes(bytes); - byteBuf.writeBytes(bytes); - return byteBuf; + @Test + public void testGet_single() throws IOException, OverCapacityException, ExecutionException, InterruptedException { + ObjectWALConfig config; + ObjectWALService wal; + List> appendCfList = new ArrayList<>(); + for (int r = 0; r < 3; r++) { + config = ObjectWALConfig.builder().withEpoch(r).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + + // append new record and verify + for (int i = 0; i < 10; i++) { + appendCfList.add(wal.append(TraceContext.DEFAULT, StreamRecordBatch.of(233L, 10, r * 10 + i, 1, generateByteBuf(256)))); + } + List> getCfList = new ArrayList<>(); + for (int i = 0; i < appendCfList.size(); i++) { + AppendResult appendRst = appendCfList.get(i).get(); + getCfList.add(wal.get(appendRst.recordOffset())); + } + for (int i = 0; i < getCfList.size(); i++) { + StreamRecordBatch record = getCfList.get(i).get(); + assertEquals(233L, record.getStreamId()); + assertEquals(i, record.getBaseOffset()); + } + + // restart wal and test get with multiple wal epoch + wal.shutdownGracefully(); + } } @Test - public void test() throws OverCapacityException, IOException { - List bufferList = new ArrayList<>(); - for (int i = 0; i < 100; i++) { - ByteBuf byteBuf = generateByteBuf(20); - bufferList.add(byteBuf); + public void testGet_batch() throws Exception { + ObjectWALConfig config; + ObjectWALService wal; + List> appendCfList = new ArrayList<>(); + for (int r = 0; r < 3; r++) { + config = ObjectWALConfig.builder().withEpoch(r).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + // append new record and verify + for (int i = 0; i < 10; i++) { + appendCfList.add(wal.append(TraceContext.DEFAULT, StreamRecordBatch.of(233L, 10, r * 10 + i, 1, generateByteBuf(256)))); + } + ((DefaultWriter) (wal.writer)).flush().join(); + for (int i = 0; i < appendCfList.size() - 3; i++) { + List records = wal.get( + DefaultRecordOffset.of(appendCfList.get(i).get().recordOffset()), + DefaultRecordOffset.of(appendCfList.get(i + 3).get().recordOffset()) + ).get(); + assertEquals(3, records.size()); + for (int j = 0; j < 3; j++) { + assertEquals(i + j, records.get(j).getBaseOffset()); + } + } + // TODO: wal end offset + + // restart wal and test get with multiple wal epoch + wal.shutdownGracefully(); } + } - List> futureList = new ArrayList<>(); - for (ByteBuf byteBuf : bufferList) { - AppendResult result = wal.append(TraceContext.DEFAULT, byteBuf.retainedSlice().asReadOnly(), 0); - futureList.add(result.future()); + @Test + public void testGet_batchSkipTrim() throws Exception { + ObjectWALConfig config = ObjectWALConfig.builder() + .withEpoch(1L) + .withMaxBytesInBatch(1024) + .withBatchInterval(1000) + .build(); + ObjectWALService wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); - if (futureList.size() == 3) { - wal.accumulator().unsafeUpload(false); - CompletableFuture.allOf(futureList.toArray(new CompletableFuture[] {})).join(); - futureList.clear(); + List> appendCfList = new ArrayList<>(); + for (int i = 0; i < 8; i++) { + appendCfList.add(wal.append(TraceContext.DEFAULT, + StreamRecordBatch.of(233L, 10, 100L + i, 1, generateByteBuf(256)))); + // ensure objects are flushed/uploaded + ((DefaultWriter) (wal.writer)).flush().join(); + if (i == 4) { + // write a trim marker by trimming to the 4th record's offset (this will produce a trim record object) + wal.trim(appendCfList.get(3).get().recordOffset()).get(); } } - if (!futureList.isEmpty()) { - wal.accumulator().unsafeUpload(false); - CompletableFuture.allOf(futureList.toArray(new CompletableFuture[] {})).join(); + + // query across a range that spans objects including the trim marker + List records = wal.get( + DefaultRecordOffset.of(appendCfList.get(4).get().recordOffset()), + DefaultRecordOffset.of(wal.confirmOffset()) + ).get(); + + assertEquals(4, records.size()); + for (int i = 0; i < records.size(); i++) { + assertEquals(104L + i, records.get(i).getBaseOffset()); + } + wal.shutdownGracefully(); + } + + + @Test + public void testTrim() throws Exception { + ObjectWALConfig config = ObjectWALConfig.builder().withEpoch(1L).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + ObjectWALService wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + + List> appendCfList = new ArrayList<>(); + for (int i = 0; i < 8; i++) { + appendCfList.add(wal.append(TraceContext.DEFAULT, StreamRecordBatch.of(233L, 0, 100L + i, 1, generateByteBuf(1)))); + if (i % 2 == 0) { + ((DefaultWriter) (wal.writer)).flush().join(); + } } - List objectList = wal.accumulator().objectList(); - assertFalse(objectList.isEmpty()); - assertTrue(objectList.size() < 100); + wal.trim(appendCfList.get(1).get().recordOffset()).get(); - // Close S3 WAL to flush all buffering data to object storage. wal.shutdownGracefully(); - // Recreate S3 WAL. - wal = new ObjectWALService(Time.SYSTEM, objectStorage, ObjectWALConfig.builder().build()); + wal = new ObjectWALService(time, objectStorage, config); wal.start(); - Iterator iterator = wal.recover(); - for (ByteBuf byteBuf : bufferList) { - assertTrue(iterator.hasNext()); + List records = new ArrayList<>(); + wal.recover().forEachRemaining(records::add); - ByteBuf recoveredByteBuf = iterator.next().record(); - assertEquals(byteBuf, recoveredByteBuf); - recoveredByteBuf.release(); + assertEquals(6, records.size()); + for (int i = 0; i < records.size(); i++) { + assertEquals(102L + i, records.get(i).record().getBaseOffset()); } - assertFalse(iterator.hasNext()); - - // Test recover after trim. - // Trim the first 2 records. - wal.trim((RecordHeader.RECORD_HEADER_SIZE + 20) * 2).join(); - assertEquals(66, wal.accumulator().objectList().size()); - - iterator = wal.recover(); - long count = 0; - while (iterator.hasNext()) { - ByteBuf record = iterator.next().record(); - record.release(); - count++; + } + + @Test + public void testRecover() throws Exception { + ObjectWALConfig config; + ObjectWALService wal; + int trimIndex = 0; + List> appendCfList = new ArrayList<>(); + for (int r = 0; r < 4; r++) { + config = ObjectWALConfig.builder().withEpoch(r).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + List records = new ArrayList<>(); + wal.recover().forEachRemaining(records::add); + // expect keep all records after trim offset + for (int i = 0; i < records.size(); i++) { + assertEquals(trimIndex + i + 1, records.get(i).record().getBaseOffset()); + } + if (r == 3) { + break; + } + for (int i = 0; i < 10; i++) { + appendCfList.add(wal.append(TraceContext.DEFAULT, StreamRecordBatch.of(233L, 10, r * 10 + i, 1, generateByteBuf(256)))); + } + ((DefaultWriter) (wal.writer)).flush().join(); + trimIndex = r * 9; + wal.trim(appendCfList.get(trimIndex).get().recordOffset()).get(); + wal.shutdownGracefully(); } - assertEquals(98, count); - - // Trim the first 3 records. - wal.trim((RecordHeader.RECORD_HEADER_SIZE + 20) * 3).join(); - assertEquals(65, wal.accumulator().objectList().size()); - - iterator = wal.recover(); - count = 0; - while (iterator.hasNext()) { - ByteBuf record = iterator.next().record(); - record.release(); - count++; + } + + @Test + public void testReset() throws Exception { + ObjectWALConfig config; + ObjectWALService wal; + int resetIndex = 0; + List> appendCfList = new ArrayList<>(); + for (int r = 0; r < 4; r++) { + config = ObjectWALConfig.builder().withEpoch(r).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + List records = new ArrayList<>(); + wal.recover().forEachRemaining(records::add); + if (r != 0) { + assertEquals(10, records.size()); + } + for (int i = 0; i < records.size(); i++) { + assertEquals(resetIndex + i, records.get(i).record().getBaseOffset()); + } + resetIndex = appendCfList.size(); + wal.reset().get(); + for (int i = 0; i < 10; i++) { + appendCfList.add(wal.append(TraceContext.DEFAULT, StreamRecordBatch.of(233L, 10, r * 10 + i, 1, generateByteBuf(256)))); + } + wal.shutdownGracefully(); } - assertEquals(97, count); + } + + public static Stream testRecoverIteratorGetContinuousFromTrimOffsetData() { + return Stream.of( + Arguments.of( + "basic", + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(20, 30)), + -1L, + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(20, 30)) + ), + Arguments.of( + "empty", + List.of(), + -1L, + List.of() + ), + Arguments.of( + "discontinuous", + List.of(mockWALObject(0, 10), mockWALObject(20, 30)), + -1L, + List.of(mockWALObject(0, 10)) + ), + Arguments.of( + "trimmed at boundary", + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(20, 30)), + 10L, + List.of(mockWALObject(10, 20), mockWALObject(20, 30)) + ), + Arguments.of( + "trimmed in middle", + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(20, 30)), + 15L, + List.of(mockWALObject(10, 20), mockWALObject(20, 30)) + ), + Arguments.of( + "trimmed nothing", + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(20, 30)), + 10L, + List.of(mockWALObject(10, 20), mockWALObject(20, 30)) + ), + Arguments.of( + "trimmed all", + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(20, 30)), + 30L, + List.of() + ), + Arguments.of( + "trimmed and discontinuous", + List.of(mockWALObject(0, 10), mockWALObject(10, 20), mockWALObject(30, 40)), + 10L, + List.of(mockWALObject(10, 20)) + ) + ); + } + + private static WALObject mockWALObject(long start, long end) { + return new WALObject((short) 0, String.format("%d-%d", start, end), 0, start, end, end - start); + } + + @ParameterizedTest(name = "Test {index} {0}") + @MethodSource("testRecoverIteratorGetContinuousFromTrimOffsetData") + public void testRecoverIteratorGetContinuousFromTrimOffset( + String name, + List objectList, + long trimOffset, + List expected + ) { + List got = getContinuousFromTrimOffset(objectList, trimOffset); + assertEquals(expected, got, name); + } + + @Test + public void testRecoverDiscontinuousObjects() throws IOException, OverCapacityException, InterruptedException, ExecutionException { + ObjectWALConfig config = ObjectWALConfig.builder().withEpoch(1L).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + ObjectWALService wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + + // write 4 objects + for (int i = 0; i < 4; i++) { + wal.append(TraceContext.DEFAULT, StreamRecordBatch.of(233L, 0, 100L + i, 1, generateByteBuf(1))); + ((DefaultWriter) (wal.writer)).flush().join(); + } + + wal.shutdownGracefully(); + + // Delete the 3nd wal object to mock it upload fail. + String nodePrefix = ObjectUtils.nodePrefix(config.clusterId(), config.nodeId()); + WALObject walObject = ObjectUtils.parse(objectStorage.list(nodePrefix).get()).get(2); + objectStorage.delete(List.of(new ObjectStorage.ObjectPath(objectStorage.bucketId(), walObject.path()))).get(); + + wal = new ObjectWALService(time, objectStorage, config); + wal.start(); + + List records = new ArrayList<>(); + wal.recover().forEachRemaining(records::add); + + assertEquals(2, records.size()); + assertEquals(100L, records.get(0).record().getBaseOffset()); + assertEquals(101L, records.get(1).record().getBaseOffset()); + } + + @Test + public void testRecoverFromV0Objects() throws IOException { + ObjectWALConfig config = ObjectWALConfig.builder().withEpoch(1L).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + + long startOffset = 0L; + for (int i = 0; i < 4; i++) { + startOffset = writeV0Object(config, StreamRecordBatch.of(233L, 0, 100L + i, 1, generateByteBuf(1)).encoded(), startOffset); + } + + ObjectWALService wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + + List records = new ArrayList<>(); + wal.recover().forEachRemaining(records::add); + + assertEquals(4, records.size()); + for (int i = 0; i < 4; i++) { + assertEquals(100L + i, records.get(i).record().getBaseOffset()); + } + } + + @Test + public void testRecoverFromV0AndV1Objects() throws IOException { + ObjectWALConfig config = ObjectWALConfig.builder().withEpoch(1L).withMaxBytesInBatch(1024).withBatchInterval(1000).build(); + long nextOffset = 0L; + nextOffset = writeV0Object(config, StreamRecordBatch.of(233L, 0, 100L, 1, generateByteBuf(1)).encoded(), nextOffset); + long record1Offset = nextOffset; + nextOffset = writeV0Object(config, StreamRecordBatch.of(233L, 0, 101L, 1, generateByteBuf(1)).encoded(), nextOffset); + nextOffset = writeV1Object(config, StreamRecordBatch.of(233L, 0, 102L, 1, generateByteBuf(1)).encoded(), nextOffset, false, 0); + nextOffset = writeV1Object(config, StreamRecordBatch.of(233L, 0, 103L, 1, generateByteBuf(1)).encoded(), nextOffset, false, record1Offset); + + ObjectWALService wal = new ObjectWALService(time, objectStorage, config); + acquire(config); + wal.start(); + List records = new ArrayList<>(); + wal.recover().forEachRemaining(records::add); + + assertEquals(2, records.size()); + for (int i = 2; i < 4; i++) { + assertEquals(100L + i, records.get(i - 2).record().getBaseOffset()); + } + } + + private long writeV0Object(ObjectWALConfig config, ByteBuf data, long startOffset) { + data = addRecordHeader(data, startOffset); + long endOffset = startOffset + data.readableBytes(); + + String path = ObjectUtils.genObjectPathV0(ObjectUtils.nodePrefix(config.clusterId(), config.nodeId()), config.epoch(), startOffset); + + CompositeByteBuf buffer = ByteBufAlloc.compositeByteBuffer(); + WALObjectHeader header = new WALObjectHeader(startOffset, data.readableBytes(), 0, 0, 0); + buffer.addComponents(true, header.marshal(), data); + + objectStorage.write(new ObjectStorage.WriteOptions(), path, buffer).join(); + return endOffset; + } + + private long writeV1Object(ObjectWALConfig config, ByteBuf data, long startOffset, boolean align, long trimOffset) { + data = addRecordHeader(data, startOffset); + long endOffset; + if (align) { + endOffset = ObjectUtils.ceilAlignOffset(startOffset); + } else { + endOffset = startOffset + data.readableBytes(); + } + String path = ObjectUtils.genObjectPathV1(ObjectUtils.nodePrefix(config.clusterId(), config.nodeId()), config.epoch(), startOffset, endOffset); + + CompositeByteBuf buffer = ByteBufAlloc.compositeByteBuffer(); + WALObjectHeader header = new WALObjectHeader(startOffset, data.readableBytes(), 0, 0, 0, trimOffset); + buffer.addComponents(true, header.marshal(), data); + + objectStorage.write(new ObjectStorage.WriteOptions(), path, buffer).join(); + return endOffset; + } + + private ByteBuf addRecordHeader(ByteBuf data, long startOffset) { + ByteBuf header = ByteBufAlloc.byteBuffer(RECORD_HEADER_SIZE); + Record record = WALUtil.generateRecord(data, header, 0, startOffset); + + CompositeByteBuf buffer = ByteBufAlloc.compositeByteBuffer(); + buffer.addComponents(true, record.header(), record.body()); + return buffer; + } + + private void acquire(ObjectWALConfig config) { + new ObjectReservationService(config.clusterId(), objectStorage, objectStorage.bucketId()) + .acquire(config.nodeId(), config.epoch(), false) + .join(); + } + + private ByteBuf generateByteBuf(int size) { + ByteBuf byteBuf = Unpooled.buffer(size); + byte[] bytes = new byte[size]; + random.nextBytes(bytes); + byteBuf.writeBytes(bytes); + return byteBuf; } } diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/RecordAccumulatorTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/RecordAccumulatorTest.java deleted file mode 100644 index bf1e9df5dc..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/RecordAccumulatorTest.java +++ /dev/null @@ -1,410 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.impl.object; - -import com.automq.stream.s3.operator.MemoryObjectStorage; -import com.automq.stream.s3.operator.ObjectStorage; -import com.automq.stream.s3.operator.ObjectStorage.ReadOptions; -import com.automq.stream.s3.wal.AppendResult; -import com.automq.stream.s3.wal.exception.OverCapacityException; -import com.automq.stream.s3.wal.exception.WALFencedException; -import com.automq.stream.utils.Time; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; - -import java.time.Duration; -import java.util.List; -import java.util.Random; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentSkipListMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.TimeUnit; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; -import io.netty.buffer.Unpooled; - -import static org.awaitility.Awaitility.await; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.api.Assertions.fail; - -public class RecordAccumulatorTest { - private RecordAccumulator recordAccumulator; - private ObjectStorage objectStorage; - private ConcurrentSkipListMap generatedByteBufMap; - private Random random; - - @BeforeEach - public void setUp() { - objectStorage = new MemoryObjectStorage(); - ObjectWALConfig config = ObjectWALConfig.builder() - .withMaxBytesInBatch(115) - .withNodeId(100) - .withEpoch(1000) - .withBatchInterval(Long.MAX_VALUE) - .withStrictBatchLimit(true) - .build(); - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, config); - recordAccumulator.start(); - generatedByteBufMap = new ConcurrentSkipListMap<>(); - random = new Random(); - } - - @AfterEach - public void tearDown() { - recordAccumulator.close(); - objectStorage.close(); - } - - private ByteBuf generateByteBuf(int size) { - ByteBuf byteBuf = Unpooled.buffer(size); - byte[] bytes = new byte[size]; - random.nextBytes(bytes); - byteBuf.writeBytes(bytes); - return byteBuf; - } - - @Test - public void testOffset() throws OverCapacityException, WALFencedException { - ByteBuf byteBuf1 = generateByteBuf(50); - CompletableFuture future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf1.readableBytes(), offset -> byteBuf1.retainedSlice().asReadOnly(), future); - assertEquals(50, recordAccumulator.nextOffset()); - - recordAccumulator.unsafeUpload(true); - long flushedOffset = future.join().flushedOffset(); - assertEquals(50, flushedOffset); - assertEquals(50, recordAccumulator.flushedOffset()); - - List objectList = recordAccumulator.objectList(); - assertEquals(1, objectList.size()); - - RecordAccumulator.WALObject object = objectList.get(0); - assertEquals(WALObjectHeader.WAL_HEADER_SIZE + 50, object.length()); - ByteBuf result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); - ByteBuf headerBuf = result.readBytes(WALObjectHeader.WAL_HEADER_SIZE); - WALObjectHeader objectHeader = WALObjectHeader.unmarshal(headerBuf); - headerBuf.release(); - assertEquals(WALObjectHeader.WAL_HEADER_MAGIC_CODE, objectHeader.magicCode()); - assertEquals(0, objectHeader.startOffset()); - assertEquals(50, objectHeader.length()); - // The last write timestamp is not set currently. - assertEquals(0L, objectHeader.stickyRecordLength()); - assertEquals(100, objectHeader.nodeId()); - assertEquals(1000, objectHeader.epoch()); - - assertEquals(byteBuf1, result); - byteBuf1.release(); - - // Test huge record. - ByteBuf byteBuf2 = generateByteBuf(50); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf2.readableBytes(), offset -> byteBuf2.retainedSlice().asReadOnly(), future); - assertEquals(100, recordAccumulator.nextOffset()); - - ByteBuf byteBuf3 = generateByteBuf(75); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf3.readableBytes(), offset -> byteBuf3.retainedSlice().asReadOnly(), future); - assertEquals(175, recordAccumulator.nextOffset()); - - recordAccumulator.unsafeUpload(true); - flushedOffset = future.join().flushedOffset(); - assertEquals(175, flushedOffset); - assertEquals(175, recordAccumulator.flushedOffset()); - - objectList = recordAccumulator.objectList(); - assertEquals(2, objectList.size()); - - object = objectList.get(1); - assertEquals(WALObjectHeader.WAL_HEADER_SIZE + 50 + 75, object.length()); - result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); - result.skipBytes(WALObjectHeader.WAL_HEADER_SIZE); - CompositeByteBuf compositeBuffer = Unpooled.compositeBuffer(); - compositeBuffer.addComponents(true, byteBuf2); - compositeBuffer.addComponents(true, byteBuf3); - assertEquals(compositeBuffer, result); - compositeBuffer.release(); - - // Test record part - ByteBuf byteBuf4 = generateByteBuf(50); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf4.readableBytes(), offset -> byteBuf4.retainedSlice().asReadOnly(), future); - assertEquals(225, recordAccumulator.nextOffset()); - - ByteBuf byteBuf5 = generateByteBuf(50); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf5.readableBytes(), offset -> byteBuf5.retainedSlice().asReadOnly(), future); - assertEquals(275, recordAccumulator.nextOffset()); - - recordAccumulator.unsafeUpload(true); - flushedOffset = future.join().flushedOffset(); - assertEquals(275, flushedOffset); - assertEquals(275, recordAccumulator.flushedOffset()); - - objectList = recordAccumulator.objectList(); - assertEquals(4, objectList.size()); - - object = objectList.get(2); - assertEquals(115, object.length()); - result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); - result.skipBytes(WALObjectHeader.WAL_HEADER_SIZE); - assertEquals(byteBuf4, result.readBytes(50)); - - object = objectList.get(3); - compositeBuffer = Unpooled.compositeBuffer(); - compositeBuffer.addComponents(true, result); - result = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); - result.skipBytes(WALObjectHeader.WAL_HEADER_SIZE); - compositeBuffer.addComponents(true, result); - assertEquals(compositeBuffer, byteBuf5); - byteBuf4.release(); - byteBuf5.release(); - } - - @Test - public void testStrictBatchLimit() throws OverCapacityException, WALFencedException { - CompletableFuture future = new CompletableFuture<>(); - recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); - recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); - recordAccumulator.append(50, offset -> generateByteBuf(50), future); - assertEquals(150, recordAccumulator.nextOffset()); - - recordAccumulator.unsafeUpload(true); - future.join(); - - assertEquals(2, recordAccumulator.objectList().size()); - - // Reset the RecordAccumulator with strict batch limit disabled. - recordAccumulator.close(); - ObjectWALConfig config = ObjectWALConfig.builder() - .withMaxBytesInBatch(115) - .withNodeId(100) - .withEpoch(1000) - .withBatchInterval(Long.MAX_VALUE) - .withStrictBatchLimit(false) - .build(); - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, config); - recordAccumulator.start(); - - assertEquals(2, recordAccumulator.objectList().size()); - - future = new CompletableFuture<>(); - recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); - recordAccumulator.append(50, offset -> generateByteBuf(50), new CompletableFuture<>()); - recordAccumulator.append(50, offset -> generateByteBuf(50), future); - assertEquals(300, recordAccumulator.nextOffset()); - - - recordAccumulator.unsafeUpload(true); - future.join(); - - assertEquals(3, recordAccumulator.objectList().size()); - } - - @ParameterizedTest - @ValueSource(booleans = {true, false}) - public void testInMultiThread(boolean strictBathLimit) throws InterruptedException, WALFencedException { - recordAccumulator.close(); - - ObjectWALConfig config = ObjectWALConfig.builder() - .withMaxBytesInBatch(115) - .withNodeId(100) - .withEpoch(1000) - .withBatchInterval(Long.MAX_VALUE) - .withStrictBatchLimit(strictBathLimit) - .build(); - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, config); - recordAccumulator.start(); - - int threadCount = 10; - CountDownLatch startBarrier = new CountDownLatch(threadCount); - CountDownLatch stopCountDownLatch = new CountDownLatch(threadCount); - List> futureList = new CopyOnWriteArrayList<>(); - for (int i = 0; i < threadCount; i++) { - new Thread(() -> { - startBarrier.countDown(); - try { - startBarrier.await(); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } - - for (int j = 0; j < 100; j++) { - ByteBuf byteBuf = generateByteBuf(40); - try { - CompletableFuture future = new CompletableFuture<>(); - long offset = recordAccumulator.append(byteBuf.readableBytes(), o -> byteBuf.retainedSlice().asReadOnly(), future); - futureList.add(future); - generatedByteBufMap.put(offset, byteBuf); - - Thread.sleep(15); - } catch (Exception e) { - throw new RuntimeException(e); - } - } - - stopCountDownLatch.countDown(); - }).start(); - } - - stopCountDownLatch.await(); - - // Ensure all records are uploaded. - try { - recordAccumulator.unsafeUpload(true); - } catch (Exception e) { - fail(e); - } - - for (CompletableFuture future : futureList) { - future.join(); - } - - assertEquals(100 * threadCount, generatedByteBufMap.size()); - - assertFalse(recordAccumulator.objectList().isEmpty()); - - CompositeByteBuf source = Unpooled.compositeBuffer(); - for (ByteBuf buffer : generatedByteBufMap.values()) { - source.addComponent(true, buffer); - } - - CompositeByteBuf result = Unpooled.compositeBuffer(); - for (RecordAccumulator.WALObject object : recordAccumulator.objectList()) { - ByteBuf buf = objectStorage.rangeRead(new ReadOptions().bucket((short) 0), object.path(), 0, object.length()).join(); - buf.skipBytes(WALObjectHeader.WAL_HEADER_SIZE); - result.addComponent(true, buf); - } - - assertEquals(source, result); - source.release(); - result.release(); - } - - @Test - public void testUploadPeriodically() throws OverCapacityException, WALFencedException { - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ObjectWALConfig.builder().build()); - recordAccumulator.start(); - - assertTrue(recordAccumulator.objectList().isEmpty()); - - ByteBuf byteBuf = generateByteBuf(25); - CompletableFuture future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf.readableBytes(), o -> byteBuf.retainedSlice().asReadOnly(), future); - - await().atMost(Duration.ofSeconds(1)).until(future::isDone); - assertEquals(1, recordAccumulator.objectList().size()); - } - - @Test - public void testShutdown() throws InterruptedException, OverCapacityException, WALFencedException { - ScheduledExecutorService executorService = recordAccumulator.executorService(); - executorService.shutdown(); - executorService.awaitTermination(10, TimeUnit.SECONDS); - - ByteBuf byteBuf = generateByteBuf(25); - CompletableFuture future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf.readableBytes(), o -> byteBuf.retainedSlice().asReadOnly(), future); - - await().during(Duration.ofSeconds(1)).atMost(Duration.ofSeconds(3)).until(() -> !future.isDone()); - assertTrue(recordAccumulator.objectList().isEmpty()); - - // Flush all data to S3 when close. - recordAccumulator.close(); - assertTrue(future.isDone()); - - - ObjectWALConfig config = ObjectWALConfig.builder() - .withMaxBytesInBatch(115) - .withNodeId(100) - .withEpoch(1000) - .withBatchInterval(Long.MAX_VALUE) - .withStrictBatchLimit(true) - .build(); - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, config); - recordAccumulator.start(); - assertEquals(1, recordAccumulator.objectList().size()); - } - - @Test - public void testTrim() throws OverCapacityException, WALFencedException { - ByteBuf byteBuf1 = generateByteBuf(50); - CompletableFuture future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf1.readableBytes(), offset -> byteBuf1.retainedSlice().asReadOnly(), future); - - ByteBuf byteBuf2 = generateByteBuf(50); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf2.readableBytes(), offset -> byteBuf2.retainedSlice().asReadOnly(), future); - - recordAccumulator.unsafeUpload(true); - long flushedOffset = future.join().flushedOffset(); - assertEquals(100, flushedOffset); - assertEquals(100, recordAccumulator.flushedOffset()); - assertEquals(2, recordAccumulator.objectList().size()); - - recordAccumulator.trim(50).join(); - assertEquals(2, recordAccumulator.objectList().size()); - - recordAccumulator.trim(100).join(); - assertEquals(0, recordAccumulator.objectList().size()); - } - - @Test - public void testReset() throws OverCapacityException, WALFencedException { - ByteBuf byteBuf1 = generateByteBuf(50); - CompletableFuture future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf1.readableBytes(), offset -> byteBuf1.retainedSlice().asReadOnly(), future); - recordAccumulator.unsafeUpload(true); - future.join(); - - ByteBuf byteBuf2 = generateByteBuf(50); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf2.readableBytes(), offset -> byteBuf2.retainedSlice().asReadOnly(), future); - recordAccumulator.unsafeUpload(true); - future.join(); - - // Close and restart with another node id. - recordAccumulator.close(); - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ObjectWALConfig.builder().withEpoch(System.currentTimeMillis()).build()); - recordAccumulator.start(); - assertEquals(0, recordAccumulator.objectList().size()); - - // Close and restart with the same node id and higher node epoch. - recordAccumulator.close(); - recordAccumulator = new RecordAccumulator(Time.SYSTEM, objectStorage, ObjectWALConfig.builder().withNodeId(100).withEpoch(System.currentTimeMillis()).build()); - recordAccumulator.start(); - assertEquals(2, recordAccumulator.objectList().size()); - - ByteBuf byteBuf3 = generateByteBuf(50); - future = new CompletableFuture<>(); - recordAccumulator.append(byteBuf3.readableBytes(), offset -> byteBuf3.retainedSlice().asReadOnly(), future); - recordAccumulator.unsafeUpload(true); - future.join(); - - List objectList = recordAccumulator.objectList(); - assertEquals(3, objectList.size()); - assertEquals(byteBuf1, objectStorage.read(new ReadOptions().bucket((short) 0), objectList.get(0).path()).join().skipBytes(WALObjectHeader.WAL_HEADER_SIZE)); - assertEquals(byteBuf2, objectStorage.read(new ReadOptions().bucket((short) 0), objectList.get(1).path()).join().skipBytes(WALObjectHeader.WAL_HEADER_SIZE)); - assertEquals(byteBuf3, objectStorage.read(new ReadOptions().bucket((short) 0), objectList.get(2).path()).join().skipBytes(WALObjectHeader.WAL_HEADER_SIZE)); - - recordAccumulator.reset().join(); - assertEquals(0, recordAccumulator.objectList().size()); - } -} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/WALObjectHeaderTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/WALObjectHeaderTest.java new file mode 100644 index 0000000000..f86605fa3f --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/WALObjectHeaderTest.java @@ -0,0 +1,35 @@ +package com.automq.stream.s3.wal.impl.object; + +import org.junit.jupiter.api.Test; + +import io.netty.buffer.ByteBuf; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class WALObjectHeaderTest { + + @Test + public void testV0() { + WALObjectHeader header = new WALObjectHeader(42L, 84L, 126L, 1, 2); + assertEquals(-1, header.trimOffset()); + assertEquals(WALObjectHeader.WAL_HEADER_MAGIC_CODE_V0, header.magicCode()); + + ByteBuf buffer = header.marshal(); + assertEquals(WALObjectHeader.WAL_HEADER_SIZE_V0, buffer.readableBytes()); + + WALObjectHeader unmarshal = WALObjectHeader.unmarshal(buffer); + assertEquals(header, unmarshal); + } + + @Test + public void testV1() { + WALObjectHeader header = new WALObjectHeader(42L, 84L, 126L, 1, 2, 168L); + assertEquals(WALObjectHeader.WAL_HEADER_MAGIC_CODE_V1, header.magicCode()); + + ByteBuf buffer = header.marshal(); + assertEquals(WALObjectHeader.WAL_HEADER_SIZE_V1, buffer.readableBytes()); + + WALObjectHeader unmarshal = WALObjectHeader.unmarshal(buffer); + assertEquals(header, unmarshal); + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/util/WALBlockDeviceChannelTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/util/WALBlockDeviceChannelTest.java deleted file mode 100644 index 4c2bdfc284..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/util/WALBlockDeviceChannelTest.java +++ /dev/null @@ -1,239 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.s3.TestUtils; -import com.automq.stream.utils.ThreadUtils; -import com.automq.stream.utils.Threads; - -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.condition.EnabledOnOs; -import org.junit.jupiter.api.condition.OS; - -import java.io.IOException; -import java.util.Optional; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ThreadLocalRandom; -import java.util.concurrent.TimeUnit; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; -import io.netty.buffer.Unpooled; - -import static com.automq.stream.s3.wal.util.WALChannelTest.TEST_BLOCK_DEVICE_KEY; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -@Tag("S3Unit") -@EnabledOnOs(OS.LINUX) -public class WALBlockDeviceChannelTest { - - static final String TEST_BLOCK_DEVICE = System.getenv(TEST_BLOCK_DEVICE_KEY); - - private String getTestPath() { - return Optional.ofNullable(TEST_BLOCK_DEVICE).orElse(TestUtils.tempFilePath()); - } - - @Test - public void testSingleThreadWriteBasic() throws IOException { - final int size = 4096 + 1; - final int count = 100; - final long capacity = WALUtil.alignLargeByBlockSize(size) * count; - - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), capacity); - channel.open(); - - for (int i = 0; i < count; i++) { - ByteBuf data = TestUtils.random(size); - long pos = WALUtil.alignLargeByBlockSize(size) * i; - writeAndFlush(channel, data, pos); - } - - channel.close(); - } - - @Test - public void testSingleThreadWriteComposite() throws IOException { - final int maxSize = 4096 * 4; - final int count = 100; - final int batch = 10; - final long capacity = WALUtil.alignLargeByBlockSize(maxSize) * count; - - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), capacity); - channel.open(); - - for (int i = 0; i < count; i += batch) { - CompositeByteBuf data = Unpooled.compositeBuffer(); - for (int j = 0; j < batch; j++) { - int size = ThreadLocalRandom.current().nextInt(1, maxSize); - data.addComponent(true, TestUtils.random(size)); - } - long pos = WALUtil.alignLargeByBlockSize(maxSize) * i; - writeAndFlush(channel, data, pos); - } - - channel.close(); - } - - @Test - public void testMultiThreadWrite() throws IOException, InterruptedException { - final int size = 4096 + 1; - final int count = 1000; - final int threads = 8; - final long capacity = WALUtil.alignLargeByBlockSize(size) * count; - - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), capacity); - channel.open(); - - ExecutorService executor = Threads.newFixedThreadPool(threads, - ThreadUtils.createThreadFactory("test-block-device-channel-write-%d", false), null); - for (int i = 0; i < count; i++) { - final int index = i; - executor.submit(() -> { - ByteBuf data = TestUtils.random(size); - long pos = WALUtil.alignLargeByBlockSize(size) * index; - try { - writeAndFlush(channel, data, pos); - } catch (IOException e) { - throw new RuntimeException(e); - } - }); - } - executor.shutdown(); - assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); - - channel.close(); - } - - @Test - public void testWriteNotAlignedBufferSize() throws IOException { - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), 1 << 20); - channel.open(); - - ByteBuf data = TestUtils.random(42); - // It's ok to do this - assertDoesNotThrow(() -> writeAndFlush(channel, data, 0)); - - channel.close(); - } - - @Test - public void testWriteNotAlignedPosition() throws IOException { - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), 1 << 20); - channel.open(); - - ByteBuf data = TestUtils.random(4096); - assertThrows(AssertionError.class, () -> writeAndFlush(channel, data, 42)); - - channel.close(); - } - - @Test - public void testWriteOutOfBound() throws IOException { - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), 4096); - channel.open(); - - ByteBuf data = TestUtils.random(4096); - assertThrows(AssertionError.class, () -> writeAndFlush(channel, data, 8192)); - - channel.close(); - } - - @Test - public void testReadBasic() throws IOException { - final int size = 4096 + 1; - final int count = 100; - final long capacity = WALUtil.alignLargeByBlockSize(size) * count; - final String path = getTestPath(); - - WALBlockDeviceChannel wChannel = new WALBlockDeviceChannel(path, capacity); - wChannel.open(); - WALBlockDeviceChannel rChannel = new WALBlockDeviceChannel(path, capacity); - rChannel.open(); - - for (int i = 0; i < count; i++) { - ByteBuf data = TestUtils.random(size); - long pos = ThreadLocalRandom.current().nextLong(0, capacity - size); - pos = WALUtil.alignSmallByBlockSize(pos); - writeAndFlush(wChannel, data, pos); - - ByteBuf buf = Unpooled.buffer(size); - int read = rChannel.read(buf, pos); - assert read == size; - assert data.equals(buf); - } - - rChannel.close(); - wChannel.close(); - } - - @Test - public void testReadInside() throws IOException { - final int size = 4096 * 4 + 1; - final int count = 100; - final long capacity = WALUtil.alignLargeByBlockSize(size) * count; - final String path = getTestPath(); - - WALBlockDeviceChannel wChannel = new WALBlockDeviceChannel(path, capacity); - wChannel.open(); - WALBlockDeviceChannel rChannel = new WALBlockDeviceChannel(path, capacity); - rChannel.open(); - - for (int i = 0; i < count; i++) { - ByteBuf data = TestUtils.random(size); - long pos = ThreadLocalRandom.current().nextLong(0, capacity - size); - pos = WALUtil.alignSmallByBlockSize(pos); - writeAndFlush(wChannel, data, pos); - - int start = ThreadLocalRandom.current().nextInt(0, size - 1); - int end = ThreadLocalRandom.current().nextInt(start + 1, size); - ByteBuf buf = Unpooled.buffer(end - start); - int read = rChannel.read(buf, pos + start); - assert read == end - start; - assert data.slice(start, end - start).equals(buf); - } - - rChannel.close(); - wChannel.close(); - } - - @Test - public void testReadNotAlignedBufferSize() throws IOException { - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), 1 << 20); - channel.open(); - - ByteBuf data = Unpooled.buffer(42); - // It's ok to do this - assertDoesNotThrow(() -> channel.read(data, 0)); - - channel.close(); - } - - @Test - public void testReadNotAlignedPosition() throws IOException { - WALBlockDeviceChannel channel = new WALBlockDeviceChannel(getTestPath(), 1 << 20); - channel.open(); - - ByteBuf data = Unpooled.buffer(4096); - // It's ok to do this - assertDoesNotThrow(() -> channel.read(data, 42)); - - channel.close(); - } - - private void writeAndFlush(WALChannel channel, ByteBuf src, long position) throws IOException { - channel.write(src, position); - channel.flush(); - } -} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/util/WALChannelTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/util/WALChannelTest.java deleted file mode 100644 index 2ba05850c1..0000000000 --- a/s3stream/src/test/java/com/automq/stream/s3/wal/util/WALChannelTest.java +++ /dev/null @@ -1,84 +0,0 @@ -/* - * Copyright 2024, AutoMQ HK Limited. - * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. - * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 - */ - -package com.automq.stream.s3.wal.util; - -import com.automq.stream.s3.TestUtils; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Tag; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.nio.ByteBuffer; - -import io.netty.buffer.ByteBuf; -import io.netty.buffer.Unpooled; - -@Tag("S3Unit") -public class WALChannelTest { - public static final String TEST_BLOCK_DEVICE_KEY = "WAL_TEST_BLOCK_DEVICE"; - - WALChannel walChannel; - - @BeforeEach - void setUp() { - walChannel = WALChannel.builder(String.format("%s/WALChannelUnitTest.data", TestUtils.tempFilePath())).direct(false).capacity(1024 * 1024 * 20).build(); - try { - walChannel.open(); - } catch (IOException e) { - throw new RuntimeException(e); - } - } - - @AfterEach - void tearDown() { - walChannel.close(); - } - - ByteBuffer createRandomTextByteBuffer(int size) { - ByteBuffer byteBuffer = ByteBuffer.allocate(size); - - for (int i = 0; i < size; i++) { - byteBuffer.put("ABCDEFGH".getBytes()[i % 8]); - } - - return byteBuffer.flip(); - } - - @Test - void testWriteAndRead() throws IOException { - ByteBuf data = TestUtils.random(1024 * 3); - for (int i = 0; i < 100; i++) { - try { - walChannel.write(data, (long) i * data.readableBytes()); - walChannel.flush(); - } catch (IOException e) { - throw new RuntimeException(e); - } - } - - final String content = "Hello World"; - walChannel.write(Unpooled.wrappedBuffer(content.getBytes()), 100); - walChannel.flush(); - - ByteBuf readBuffer = Unpooled.buffer(content.length()); - int read = walChannel.read(readBuffer, 100); - - String readString = new String(readBuffer.array()); - System.out.println(new String(readBuffer.array())); - System.out.println(read); - - assert read == content.length(); - assert readString.equals(content); - } -} diff --git a/s3stream/src/test/java/com/automq/stream/utils/AbstractOrderedCollectionTest.java b/s3stream/src/test/java/com/automq/stream/utils/AbstractOrderedCollectionTest.java index 0116b476ff..0cb3af8020 100644 --- a/s3stream/src/test/java/com/automq/stream/utils/AbstractOrderedCollectionTest.java +++ b/s3stream/src/test/java/com/automq/stream/utils/AbstractOrderedCollectionTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/test/java/com/automq/stream/utils/AsyncSemaphoreTest.java b/s3stream/src/test/java/com/automq/stream/utils/AsyncSemaphoreTest.java new file mode 100644 index 0000000000..3fc0b299ac --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/utils/AsyncSemaphoreTest.java @@ -0,0 +1,228 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.automq.stream.utils; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class AsyncSemaphoreTest { + + private AsyncSemaphore semaphore; + + @BeforeEach + void setup() { + semaphore = new AsyncSemaphore(10); // Initialize with 10 permits + } + + @Test + public void testAcquireWithSufficientPermits() throws ExecutionException, InterruptedException, TimeoutException { + // Test that permits can be acquired when sufficient permits are available + AtomicInteger executedTaskNum = new AtomicInteger(0); + + CompletableFuture task1 = new CompletableFuture<>(); + boolean acquired1 = semaphore.acquire(3, () -> { + executedTaskNum.addAndGet(1); + return task1; + }, Runnable::run); + assertTrue(acquired1); // Task can acquire permits + assertEquals(7, semaphore.permits()); // Permits should be decreased + assertEquals(1, executedTaskNum.get()); // Task should execute + assertFalse(semaphore.requiredRelease()); // No release required + + CompletableFuture task2 = new CompletableFuture<>(); + boolean acquired2 = semaphore.acquire(3, () -> { + executedTaskNum.addAndGet(1); + return task2; + }, Runnable::run); + assertTrue(acquired2); // Task can acquire permits + assertEquals(4, semaphore.permits()); // Permits should be decreased + assertEquals(2, executedTaskNum.get()); // Task should execute + assertFalse(semaphore.requiredRelease()); // No release required + + CompletableFuture task3 = new CompletableFuture<>(); + boolean acquired3 = semaphore.acquire(4, () -> { + executedTaskNum.addAndGet(1); + return task3; + }, Runnable::run); + assertTrue(acquired3); // Task can acquire permits + assertEquals(0, semaphore.permits()); // Permits should be decreased + assertEquals(3, executedTaskNum.get()); // Task should execute + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + // Release permits and ensure they are restored + task1.complete(null); // Release permits from task1 + assertEquals(3, semaphore.permits()); + task2.complete(null); // Release permits from task2 + assertEquals(6, semaphore.permits()); + task3.complete(null); // Release permits from task3 + assertEquals(10, semaphore.permits()); + } + + @Test + public void testAcquireDecreaseToNegativePermits() throws ExecutionException, InterruptedException, TimeoutException { + // Test that permits can decrease to negative values + AtomicInteger executedTaskNum = new AtomicInteger(0); + + CompletableFuture task1 = new CompletableFuture<>(); + boolean acquired1 = semaphore.acquire(15, () -> { + executedTaskNum.addAndGet(1); + return task1; + }, Runnable::run); + assertTrue(acquired1); // Task can acquire permits + assertEquals(-5, semaphore.permits()); // Permits should be negative + assertEquals(1, executedTaskNum.get()); // Task should execute + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + // Test that a second task is queued when permits are negative + CompletableFuture task2 = new CompletableFuture<>(); + boolean acquired2 = semaphore.acquire(5, () -> { + executedTaskNum.addAndGet(1); + return task2; + }, Runnable::run); + assertFalse(acquired2); // Task should be queued + assertEquals(-5, semaphore.permits()); // Permits should remain unchanged + assertEquals(1, executedTaskNum.get()); // Task should not execute yet + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + // Release permits and ensure the queued task executes + task1.complete(null); // Release permits from task1 + assertEquals(5, semaphore.permits()); // Permits should be restored and acquired by task2 + assertEquals(2, executedTaskNum.get()); // Task2 should execute + assertFalse(semaphore.requiredRelease()); // No release required + + task2.complete(null); // Release permits from task2 + assertEquals(10, semaphore.permits()); // Permits should be restored + } + + @Test + public void testReleaseWithoutLeaks() throws ExecutionException, InterruptedException, TimeoutException { + // Test that all releases are correctly accounted for, even if there are exceptions when acquiring + AtomicInteger executedTaskNum = new AtomicInteger(0); + + CompletableFuture task1 = new CompletableFuture<>(); + boolean acquired1 = semaphore.acquire(5, () -> { + executedTaskNum.addAndGet(1); + return task1; + }, Runnable::run); + assertTrue(acquired1); // Task can acquire permits + assertEquals(5, semaphore.permits()); // Permits should be decreased + assertEquals(1, executedTaskNum.get()); // Task should execute + assertFalse(semaphore.requiredRelease()); // No release required + + CompletableFuture task2 = new CompletableFuture<>(); + boolean acquired2 = semaphore.acquire(10, () -> { + executedTaskNum.addAndGet(1); + return task2; + }, Runnable::run); + assertTrue(acquired2); // Task can acquire permits + assertEquals(-5, semaphore.permits()); // Permits should be negative + assertEquals(2, executedTaskNum.get()); // Task should execute + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + RuntimeException task3Exception = new RuntimeException("Task 3 exception"); + CompletableFuture task3 = new CompletableFuture<>(); + boolean acquired3 = semaphore.acquire(1, () -> { + executedTaskNum.addAndGet(1); + throw task3Exception; + }, Runnable::run); + assertFalse(acquired3); // Task should be queued + assertEquals(-5, semaphore.permits()); // Permits should remain unchanged + assertEquals(2, executedTaskNum.get()); // Task should not execute yet + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + // Release permits from task1 + task1.complete(null); + assertEquals(0, semaphore.permits()); // Permits should be restored and not acquired by task3 + assertEquals(2, executedTaskNum.get()); // Task3 should not execute + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + // Release permits from task2 + task2.complete(null); + // Permits should be restored and acquired by task3, but task3 throw an exception, so the permits should be restored + assertEquals(10, semaphore.permits()); + assertEquals(3, executedTaskNum.get()); // Task3 should execute + assertFalse(semaphore.requiredRelease()); // No release required + + // Release permits from task3 + task3.completeExceptionally(task3Exception); + assertEquals(10, semaphore.permits()); // Permits should be restored even if task3 completes exceptionally + } + + @Test + public void testSequentialAcquireOrder() throws ExecutionException, InterruptedException, TimeoutException { + // Test that sequentially issued acquire requests are executed in order after permits become sufficient + AtomicInteger executionOrder = new AtomicInteger(0); + + CompletableFuture task1 = new CompletableFuture<>(); + boolean acquired1 = semaphore.acquire(12, () -> { + executionOrder.compareAndSet(0, 1); + return task1; + }, Runnable::run); + assertTrue(acquired1); // Task 1 can acquire permits + assertEquals(-2, semaphore.permits()); // Permits should be negative + assertEquals(1, executionOrder.get()); // Task 1 should execute + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + CompletableFuture task2 = new CompletableFuture<>(); + boolean acquired2 = semaphore.acquire(6, () -> { + executionOrder.compareAndSet(1, 2); + return task2; + }, Runnable::run); + assertFalse(acquired2); // Task 2 should be queued + assertEquals(-2, semaphore.permits()); // Permits should remain unchanged + assertEquals(1, executionOrder.get()); // Task 2 should not execute yet + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + CompletableFuture task3 = new CompletableFuture<>(); + boolean acquired3 = semaphore.acquire(8, () -> { + executionOrder.compareAndSet(2, 3); + return task3; + }, Runnable::run); + assertFalse(acquired3); // Task 3 should be queued + assertEquals(-2, semaphore.permits()); // Permits should remain unchanged + assertEquals(1, executionOrder.get()); // Task 3 should not execute yet + assertTrue(semaphore.requiredRelease()); // Release required due to non-positive permits + + // Release permits from task1 and ensure task2 executes first + task1.complete(null); + assertEquals(4, semaphore.permits()); // Permits should be restored and acquired by task2 + assertEquals(2, executionOrder.get()); // Task 2 should execute first + assertTrue(semaphore.requiredRelease()); // Release required due to non-empty queue + + // Release permits from task2 and ensure task3 executes next + task2.complete(null); + assertEquals(2, semaphore.permits()); // Permits should be restored and acquired by task3 + assertEquals(3, executionOrder.get()); // Task 3 should execute next + assertFalse(semaphore.requiredRelease()); // No release required + + // Release permits from task3 + task3.complete(null); + assertEquals(10, semaphore.permits()); // Permits should be restored + } +} diff --git a/s3stream/src/test/java/com/automq/stream/utils/FutureTickerTest.java b/s3stream/src/test/java/com/automq/stream/utils/FutureTickerTest.java index 4b11ee3e90..abc8d53c0b 100644 --- a/s3stream/src/test/java/com/automq/stream/utils/FutureTickerTest.java +++ b/s3stream/src/test/java/com/automq/stream/utils/FutureTickerTest.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package com.automq.stream.utils; diff --git a/s3stream/src/test/resources/log4j.properties b/s3stream/src/test/resources/log4j.properties deleted file mode 100644 index 5da498bd9c..0000000000 --- a/s3stream/src/test/resources/log4j.properties +++ /dev/null @@ -1,15 +0,0 @@ -# -# Copyright 2024, AutoMQ HK Limited. -# -# The use of this file is governed by the Business Source License, -# as detailed in the file "/LICENSE.S3Stream" included in this repository. -# -# As of the Change Date specified in that file, in accordance with -# the Business Source License, use of this software will be governed -# by the Apache License, Version 2.0 -# -log4j.rootLogger=OFF, stdout -log4j.appender.stdout=org.apache.log4j.ConsoleAppender -log4j.appender.stdout.layout=org.apache.log4j.PatternLayout -log4j.appender.stdout.layout.ConversionPattern=[%d] %p %m (%c:%L)%n -log4j.logger.com.automq=WARN diff --git a/server-common/src/main/java/org/apache/kafka/common/WeightedRandomList.java b/server-common/src/main/java/org/apache/kafka/common/WeightedRandomList.java index ef39401f50..487bf810a7 100644 --- a/server-common/src/main/java/org/apache/kafka/common/WeightedRandomList.java +++ b/server-common/src/main/java/org/apache/kafka/common/WeightedRandomList.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.common; diff --git a/server-common/src/main/java/org/apache/kafka/server/common/automq/AutoMQVersion.java b/server-common/src/main/java/org/apache/kafka/server/common/automq/AutoMQVersion.java index 3f0c01d3cf..d56c124d1c 100644 --- a/server-common/src/main/java/org/apache/kafka/server/common/automq/AutoMQVersion.java +++ b/server-common/src/main/java/org/apache/kafka/server/common/automq/AutoMQVersion.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.server.common.automq; @@ -23,10 +31,12 @@ public enum AutoMQVersion { // Support object bucket index // Support huge cluster // Support node registration - V2((short) 3); + V2((short) 3), + // Support zero zone v2 + V3((short) 4); public static final String FEATURE_NAME = "automq.version"; - public static final AutoMQVersion LATEST = V2; + public static final AutoMQVersion LATEST = V3; private final short level; private final Version s3streamVersion; @@ -85,6 +95,10 @@ public boolean isNodeRegistrationSupported() { return isAtLeast(V2); } + public boolean isZeroZoneV2Supported() { + return isAtLeast(V3); + } + public short streamRecordVersion() { if (isReassignmentV1Supported()) { return 1; @@ -131,6 +145,7 @@ private Version mapS3StreamVersion(short automqVersion) { case 2: return Version.V0; case 3: + case 4: return Version.V1; default: throw new IllegalArgumentException("Unknown AutoMQVersion level: " + automqVersion); diff --git a/server-common/src/main/java/org/apache/kafka/server/common/automq/TableTopicConfigValidator.java b/server-common/src/main/java/org/apache/kafka/server/common/automq/TableTopicConfigValidator.java new file mode 100644 index 0000000000..92be852a38 --- /dev/null +++ b/server-common/src/main/java/org/apache/kafka/server/common/automq/TableTopicConfigValidator.java @@ -0,0 +1,167 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.server.common.automq; + +import org.apache.kafka.common.config.ConfigDef; +import org.apache.kafka.common.config.ConfigException; + +import org.apache.commons.lang3.tuple.Pair; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import static java.util.stream.Collectors.toList; + +public class TableTopicConfigValidator { + static final Pattern LIST_STRING_REGEX = Pattern.compile("\\[(.*)\\]"); + public static final String COMMA_NO_PARENS_REGEX = ",(?![^()]*+\\))"; + private static final Pattern COLUMN_NAME_REGEX = Pattern.compile("[a-zA-Z0-9_\\-\\.]+"); + + public static class PartitionValidator implements ConfigDef.Validator { + public static final PartitionValidator INSTANCE = new PartitionValidator(); + private static final Pattern TRANSFORM_REGEX = Pattern.compile("(\\w+)\\((.+)\\)"); + + @Override + public void ensureValid(String name, Object value) { + if (value == null) { + return; + } + try { + String str = (String) value; + List partitions = parsePartitionBy(str); + partitions.forEach(partitionField -> { + Matcher matcher = TRANSFORM_REGEX.matcher(partitionField); + if (matcher.matches()) { + Transform transform = Transform.fromString(matcher.group(1)); + switch (transform) { + case YEAR: + case MONTH: + case DAY: + case HOUR: + String columnName = matcher.group(2); + if (!COLUMN_NAME_REGEX.matcher(columnName).matches()) { + throw new ConfigException(name, value, String.format("Invalid column name %s", columnName)); + } + break; + case BUCKET: + case TRUNCATE: + Pair pair = transformArgPair(matcher.group(2)); + if (!COLUMN_NAME_REGEX.matcher(pair.getLeft()).matches()) { + throw new ConfigException(name, value, String.format("Invalid column name %s", pair.getLeft())); + } + break; + default: + throw new IllegalArgumentException("Unsupported transform: " + transform); + } + } else { + if (!COLUMN_NAME_REGEX.matcher(partitionField).matches()) { + throw new ConfigException(name, value, String.format("Invalid column name %s", partitionField)); + } + } + }); + + } catch (Throwable e) { + if (e instanceof ConfigException) { + throw e; + } else { + throw new ConfigException(name, value, e.getMessage()); + } + } + + } + + public static List parsePartitionBy(String str) { + return stringToList(str, COMMA_NO_PARENS_REGEX); + } + + public static Pair transformArgPair(String argsStr) { + String[] parts = argsStr.split(","); + if (parts.length != 2) { + throw new IllegalArgumentException("Invalid argument " + argsStr + ", should have 2 parts"); + } + return Pair.of(parts[0].trim(), Integer.parseInt(parts[1].trim())); + } + } + + public static class IdColumnsValidator implements ConfigDef.Validator { + public static final IdColumnsValidator INSTANCE = new IdColumnsValidator(); + + @Override + public void ensureValid(String name, Object value) { + if (value == null) { + return; + } + try { + String str = (String) value; + List idColumns = stringToList(str, COMMA_NO_PARENS_REGEX); + idColumns.forEach(idColumn -> { + if (!COLUMN_NAME_REGEX.matcher(idColumn).matches()) { + throw new ConfigException(name, value, String.format("Invalid column name %s", idColumn)); + } + }); + } catch (Throwable e) { + if (e instanceof ConfigException) { + throw e; + } else { + throw new ConfigException(name, value, e.getMessage()); + } + } + } + } + + public static List stringToList(String value, String regex) { + if (value == null || value.isEmpty()) { + return Collections.emptyList(); + } + Matcher matcher = LIST_STRING_REGEX.matcher(value); + if (matcher.matches()) { + value = matcher.group(1); + } + return Arrays.stream(value.split(regex)).map(String::trim).collect(toList()); + } + + public enum Transform { + YEAR, MONTH, DAY, HOUR, BUCKET, TRUNCATE; + + public static Transform fromString(String str) { + switch (str.toLowerCase(Locale.ROOT)) { + case "year": + return YEAR; + case "month": + return MONTH; + case "day": + return DAY; + case "hour": + return HOUR; + case "bucket": + return BUCKET; + case "truncate": + return TRUNCATE; + default: + throw new IllegalArgumentException("Invalid transform function " + str); + } + } + } + +} diff --git a/server-common/src/main/java/org/apache/kafka/server/config/QuotaConfigs.java b/server-common/src/main/java/org/apache/kafka/server/config/QuotaConfigs.java index db5ecdd872..b1e5e2a020 100644 --- a/server-common/src/main/java/org/apache/kafka/server/config/QuotaConfigs.java +++ b/server-common/src/main/java/org/apache/kafka/server/config/QuotaConfigs.java @@ -104,9 +104,15 @@ public class QuotaConfigs { "for the specified IP."; // AutoMQ inject start + /** + * All clients created by AutoMQ will have this prefix in their client id, and they will be excluded from quota. + */ + public static final String INTERNAL_CLIENT_ID_PREFIX = "__automq_client_"; + public static final String BROKER_QUOTA_ENABLED_CONFIG = "broker.quota.enabled"; public static final String BROKER_QUOTA_PRODUCE_BYTES_CONFIG = "broker.quota.produce.bytes"; public static final String BROKER_QUOTA_FETCH_BYTES_CONFIG = "broker.quota.fetch.bytes"; + public static final String BROKER_QUOTA_SLOW_FETCH_BYTES_CONFIG = "broker.quota.slow.fetch.bytes"; public static final String BROKER_QUOTA_REQUEST_RATE_CONFIG = "broker.quota.request.rate"; public static final String BROKER_QUOTA_WHITE_LIST_USER_CONFIG = "broker.quota.white.list.user"; public static final String BROKER_QUOTA_WHITE_LIST_CLIENT_ID_CONFIG = "broker.quota.white.list.client.id"; @@ -117,6 +123,7 @@ public class QuotaConfigs { public static final String BROKER_QUOTA_ENABLED_DOC = "Enable broker quota."; public static final String BROKER_QUOTA_PRODUCE_BYTES_DOC = "The maximum bytes send by producer in single window."; public static final String BROKER_QUOTA_FETCH_BYTES_DOC = "The maximum bytes receive by consumer in single window."; + public static final String BROKER_QUOTA_SLOW_FETCH_BYTES_DOC = "The maximum bytes receive by slow fetch consumer in single window."; public static final String BROKER_QUOTA_REQUEST_RATE_DOC = "The maximum request count send by client in single window."; public static final String BROKER_QUOTA_WHITE_LIST_USER_DOC = "Broker quota white list for user."; public static final String BROKER_QUOTA_WHITE_LIST_CLIENT_ID_DOC = "Broker quota white list for client id."; @@ -176,6 +183,7 @@ public static ConfigDef brokerQuotaConfigs() { .define(QuotaConfigs.BROKER_QUOTA_ENABLED_CONFIG, ConfigDef.Type.BOOLEAN, false, MEDIUM, QuotaConfigs.BROKER_QUOTA_ENABLED_DOC) .define(QuotaConfigs.BROKER_QUOTA_PRODUCE_BYTES_CONFIG, ConfigDef.Type.DOUBLE, Double.MAX_VALUE, MEDIUM, QuotaConfigs.BROKER_QUOTA_PRODUCE_BYTES_DOC) .define(QuotaConfigs.BROKER_QUOTA_FETCH_BYTES_CONFIG, ConfigDef.Type.DOUBLE, Double.MAX_VALUE, MEDIUM, QuotaConfigs.BROKER_QUOTA_FETCH_BYTES_DOC) + .define(QuotaConfigs.BROKER_QUOTA_SLOW_FETCH_BYTES_CONFIG, ConfigDef.Type.DOUBLE, Double.MAX_VALUE, MEDIUM, QuotaConfigs.BROKER_QUOTA_SLOW_FETCH_BYTES_DOC) .define(QuotaConfigs.BROKER_QUOTA_REQUEST_RATE_CONFIG, ConfigDef.Type.DOUBLE, Double.MAX_VALUE, MEDIUM, QuotaConfigs.BROKER_QUOTA_REQUEST_RATE_DOC) .define(QuotaConfigs.BROKER_QUOTA_WHITE_LIST_USER_CONFIG, ConfigDef.Type.STRING, "", MEDIUM, QuotaConfigs.BROKER_QUOTA_WHITE_LIST_USER_DOC) .define(QuotaConfigs.BROKER_QUOTA_WHITE_LIST_CLIENT_ID_CONFIG, ConfigDef.Type.STRING, "", MEDIUM, QuotaConfigs.BROKER_QUOTA_WHITE_LIST_CLIENT_ID_DOC) diff --git a/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java b/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java index 66747e7436..ed30f872c0 100644 --- a/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java +++ b/server-common/src/main/java/org/apache/kafka/server/config/ServerTopicConfigSynonyms.java @@ -83,6 +83,18 @@ public final class ServerTopicConfigSynonyms { sameName(TopicConfig.COMPRESSION_GZIP_LEVEL_CONFIG), sameName(TopicConfig.COMPRESSION_LZ4_LEVEL_CONFIG), sameName(TopicConfig.COMPRESSION_ZSTD_LEVEL_CONFIG), + + // AutoMQ inject start + sameName(TopicConfig.TABLE_TOPIC_ENABLE_CONFIG), + sameName(TopicConfig.TABLE_TOPIC_COMMIT_INTERVAL_CONFIG), + sameName(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG), + sameName(TopicConfig.TABLE_TOPIC_SCHEMA_TYPE_CONFIG), + + sameName(TopicConfig.KAFKA_LINKS_ID_CONFIG), + sameName(TopicConfig.KAFKA_LINKS_TOPIC_START_TIME_CONFIG), + sameName(TopicConfig.KAFKA_LINKS_TOPIC_STATE_CONFIG), + // AutoMQ inject end + sameNameWithLogPrefix(TopicConfig.PREALLOCATE_CONFIG), sameNameWithLogPrefix(TopicConfig.MESSAGE_FORMAT_VERSION_CONFIG), sameNameWithLogPrefix(TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG), diff --git a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/MultiAttributes.java b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/MultiAttributes.java index ca4b5206e0..087e3779d1 100644 --- a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/MultiAttributes.java +++ b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/MultiAttributes.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.server.metrics.s3stream; diff --git a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/PartitionCountDistribution.java b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/PartitionCountDistribution.java index c3b6e66228..976fa67445 100644 --- a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/PartitionCountDistribution.java +++ b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/PartitionCountDistribution.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ CO.,LTD. + * Copyright 2025, AutoMQ HK Limited. * - * Use of this software is governed by the Business Source License - * included in the file BSL.md + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.server.metrics.s3stream; diff --git a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsConstants.java b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsConstants.java index 98ea1b5bf6..45ceea2b95 100644 --- a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsConstants.java +++ b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsConstants.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.server.metrics.s3stream; @@ -20,7 +28,10 @@ public class S3StreamKafkaMetricsConstants { public static final String STREAM_SET_OBJECT_NUM = "stream_set_object_num"; public static final String STREAM_OBJECT_NUM = "stream_object_num"; public static final String FETCH_LIMITER_PERMIT_NUM = "fetch_limiter_permit_num"; + public static final String FETCH_LIMITER_WAITING_TASK_NUM = "fetch_limiter_waiting_task_num"; public static final String FETCH_PENDING_TASK_NUM = "fetch_pending_task_num"; + public static final String FETCH_LIMITER_TIMEOUT_COUNT = "fetch_limiter_timeout_count"; + public static final String FETCH_LIMITER_TIME = "fetch_limiter_time"; public static final String LOG_APPEND_PERMIT_NUM = "log_append_permit_num"; public static final String SLOW_BROKER_METRIC_NAME = "slow_broker_count"; public static final String TOPIC_PARTITION_COUNT_METRIC_NAME = "topic_partition_count"; @@ -45,4 +56,8 @@ public class S3StreamKafkaMetricsConstants { public static final String PARTITION_STATUS_STATISTICS_METRIC_NAME = "partition_status_statistics"; public static final AttributeKey LABEL_STATUS = AttributeKey.stringKey("status"); + + // Back Pressure + public static final String BACK_PRESSURE_STATE_METRIC_NAME = "back_pressure_state"; + public static final AttributeKey LABEL_BACK_PRESSURE_STATE = AttributeKey.stringKey("state"); } diff --git a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManager.java b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManager.java index 27d2426eae..084e4ed707 100644 --- a/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManager.java +++ b/server-common/src/main/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManager.java @@ -1,35 +1,63 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.server.metrics.s3stream; import com.automq.stream.s3.metrics.MetricsConfig; import com.automq.stream.s3.metrics.MetricsLevel; +import com.automq.stream.s3.metrics.NoopLongCounter; import com.automq.stream.s3.metrics.NoopObservableLongGauge; import com.automq.stream.s3.metrics.wrapper.ConfigListener; +import com.automq.stream.s3.metrics.wrapper.CounterMetric; +import com.automq.stream.s3.metrics.wrapper.HistogramInstrument; +import com.automq.stream.s3.metrics.wrapper.HistogramMetric; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import java.io.ByteArrayInputStream; +import java.security.cert.CertificateException; +import java.security.cert.CertificateFactory; +import java.security.cert.X509Certificate; import java.util.ArrayList; +import java.util.Base64; import java.util.Collections; +import java.util.Date; import java.util.List; import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.function.Function; import java.util.function.Supplier; import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.metrics.LongCounter; import io.opentelemetry.api.metrics.Meter; import io.opentelemetry.api.metrics.ObservableLongGauge; public class S3StreamKafkaMetricsManager { + private static final Logger LOGGER = LoggerFactory.getLogger(S3StreamKafkaMetricsManager.class); + private static final List BASE_ATTRIBUTES_LISTENERS = new ArrayList<>(); + + public static final List FETCH_LIMITER_TIME_METRICS = new CopyOnWriteArrayList<>(); + private static final MultiAttributes BROKER_ATTRIBUTES = new MultiAttributes<>(Attributes.empty(), S3StreamKafkaMetricsConstants.LABEL_NODE_ID); private static final MultiAttributes S3_OBJECT_ATTRIBUTES = new MultiAttributes<>(Attributes.empty(), @@ -40,12 +68,18 @@ public class S3StreamKafkaMetricsManager { S3StreamKafkaMetricsConstants.LABEL_FETCH_EXECUTOR_NAME); private static final MultiAttributes PARTITION_STATUS_STATISTICS_ATTRIBUTES = new MultiAttributes<>(Attributes.empty(), S3StreamKafkaMetricsConstants.LABEL_STATUS); + private static final MultiAttributes BACK_PRESSURE_STATE_ATTRIBUTES = new MultiAttributes<>(Attributes.empty(), + S3StreamKafkaMetricsConstants.LABEL_BACK_PRESSURE_STATE); + + // List to store all the observable long gauges for certificates + private static final List CERT_OBSERVABLE_LONG_GAUGES = new ArrayList<>(); static { BASE_ATTRIBUTES_LISTENERS.add(BROKER_ATTRIBUTES); BASE_ATTRIBUTES_LISTENERS.add(S3_OBJECT_ATTRIBUTES); BASE_ATTRIBUTES_LISTENERS.add(FETCH_LIMITER_ATTRIBUTES); BASE_ATTRIBUTES_LISTENERS.add(FETCH_EXECUTOR_ATTRIBUTES); + BASE_ATTRIBUTES_LISTENERS.add(BACK_PRESSURE_STATE_ATTRIBUTES); } private static Supplier isActiveSupplier = () -> false; @@ -59,10 +93,16 @@ public class S3StreamKafkaMetricsManager { private static Supplier> streamSetObjectNumSupplier = Collections::emptyMap; private static ObservableLongGauge streamObjectNumMetrics = new NoopObservableLongGauge(); private static Supplier streamObjectNumSupplier = () -> 0; + private static ObservableLongGauge fetchLimiterPermitNumMetrics = new NoopObservableLongGauge(); private static Supplier> fetchLimiterPermitNumSupplier = Collections::emptyMap; + private static ObservableLongGauge fetchLimiterWaitingTaskNumMetrics = new NoopObservableLongGauge(); + private static Supplier> fetchLimiterWaitingTaskNumSupplier = Collections::emptyMap; private static ObservableLongGauge fetchPendingTaskNumMetrics = new NoopObservableLongGauge(); private static Supplier> fetchPendingTaskNumSupplier = Collections::emptyMap; + private static LongCounter fetchLimiterTimeoutCount = new NoopLongCounter(); + private static HistogramInstrument fetchLimiterTime; + private static ObservableLongGauge logAppendPermitNumMetrics = new NoopObservableLongGauge(); private static Supplier logAppendPermitNumSupplier = () -> 0; private static MetricsConfig metricsConfig = new MetricsConfig(MetricsLevel.INFO, Attributes.empty()); @@ -75,6 +115,22 @@ public class S3StreamKafkaMetricsManager { private static List partitionStatusList = Collections.emptyList(); private static Function partitionStatusStatisticsSupplier = s -> 0; + private static ObservableLongGauge backPressureState = new NoopObservableLongGauge(); + /** + * Supplier for back pressure state. + * Key is the state name, value is 1 for current state, -1 for other states. + */ + private static Supplier> backPressureStateSupplier = Collections::emptyMap; + + /** + * supplier for truststoreCerts + */ + private static Supplier truststoreCertsSupplier = () -> null; + /** + * supplier for server cert chain + */ + private static Supplier certChainSupplier = () -> null; + public static void configure(MetricsConfig metricsConfig) { synchronized (BASE_ATTRIBUTES_LISTENERS) { S3StreamKafkaMetricsManager.metricsConfig = metricsConfig; @@ -90,6 +146,12 @@ public static void initMetrics(Meter meter, String prefix) { initFetchMetrics(meter, prefix); initLogAppendMetrics(meter, prefix); initPartitionStatusStatisticsMetrics(meter, prefix); + initBackPressureMetrics(meter, prefix); + try { + initCertMetrics(meter, prefix); + } catch (Exception e) { + LOGGER.error("Failed to init cert metrics", e); + } } private static void initAutoBalancerMetrics(Meter meter, String prefix) { @@ -193,6 +255,18 @@ private static void initFetchMetrics(Meter meter, String prefix) { } } }); + fetchLimiterWaitingTaskNumMetrics = meter.gaugeBuilder(prefix + S3StreamKafkaMetricsConstants.FETCH_LIMITER_WAITING_TASK_NUM) + .setDescription("The number of tasks waiting for permits in fetch limiters") + .ofLongs() + .buildWithCallback(result -> { + if (MetricsLevel.INFO.isWithin(metricsConfig.getMetricsLevel())) { + Map fetchLimiterWaitingTaskNumMap = fetchLimiterWaitingTaskNumSupplier.get(); + for (Map.Entry entry : fetchLimiterWaitingTaskNumMap.entrySet()) { + result.record(entry.getValue(), FETCH_LIMITER_ATTRIBUTES.get(entry.getKey())); + } + } + }); + fetchPendingTaskNumMetrics = meter.gaugeBuilder(prefix + S3StreamKafkaMetricsConstants.FETCH_PENDING_TASK_NUM) .setDescription("The number of pending tasks in fetch executors") .ofLongs() @@ -204,6 +278,12 @@ private static void initFetchMetrics(Meter meter, String prefix) { } } }); + + fetchLimiterTimeoutCount = meter.counterBuilder(prefix + S3StreamKafkaMetricsConstants.FETCH_LIMITER_TIMEOUT_COUNT) + .setDescription("The number of acquire permits timeout in fetch limiters") + .build(); + fetchLimiterTime = new HistogramInstrument(meter, prefix + S3StreamKafkaMetricsConstants.FETCH_LIMITER_TIME, + "The time cost of acquire permits in fetch limiters", "nanoseconds", () -> FETCH_LIMITER_TIME_METRICS); } private static void initLogAppendMetrics(Meter meter, String prefix) { @@ -230,6 +310,131 @@ private static void initPartitionStatusStatisticsMetrics(Meter meter, String pre }); } + private static void initBackPressureMetrics(Meter meter, String prefix) { + backPressureState = meter.gaugeBuilder(prefix + S3StreamKafkaMetricsConstants.BACK_PRESSURE_STATE_METRIC_NAME) + .setDescription("Back pressure state") + .ofLongs() + .buildWithCallback(result -> { + if (MetricsLevel.INFO.isWithin(metricsConfig.getMetricsLevel())) { + Map states = backPressureStateSupplier.get(); + states.forEach((state, value) -> { + result.record(value, BACK_PRESSURE_STATE_ATTRIBUTES.get(state)); + }); + } + }); + } + + /** + * Initialize the certificate metrics. + * + * @param meter The OpenTelemetry meter to use for creating metrics. + */ + public static void initCertMetrics(Meter meter, String prefix) throws CertificateException { + String truststoreCerts = truststoreCertsSupplier.get(); + String certChain = certChainSupplier.get(); + if (truststoreCerts == null || truststoreCerts.isEmpty()) { + return; + } + if (certChain == null || certChain.isEmpty()) { + return; + } + // Add TLS certificate metrics + addTlsMetrics(certChain, truststoreCerts, meter, prefix); + } + + /** + * Add TLS certificate metrics. + * + * @param certChain The certificate chain in PEM format. + * @param truststoreCerts The truststore certificates in PEM format. + * @param meter The OpenTelemetry meter to use for creating metrics. + * @param prefix The prefix for the metric names. + */ + private static void addTlsMetrics(String certChain, String truststoreCerts, Meter meter, String prefix) throws CertificateException { + // Parse and check the certificate expiration time + X509Certificate[] serverCerts = parseCertificates(certChain); + X509Certificate[] trustStoreCerts = parseCertificates(truststoreCerts); + + for (X509Certificate cert : serverCerts) { + registerCertMetrics(meter, cert, "server_cert", prefix); + } + for (X509Certificate cert : trustStoreCerts) { + registerCertMetrics(meter, cert, "truststore_cert", prefix); + } + } + + /** + * Register certificate metrics. + * + * @param meter The OpenTelemetry meter to use for creating metrics. + * @param cert The X509 certificate to register metrics for. + * @param certType The type of the certificate (e.g., "server_cert", "truststore_cert"). + * @param prefix The prefix for the metric names. + */ + private static void registerCertMetrics(Meter meter, X509Certificate cert, String certType, String prefix) { + String subject = cert.getSubjectX500Principal().getName(); + Date expiryDate = cert.getNotAfter(); + long daysRemaining = (expiryDate.getTime() - System.currentTimeMillis()) / (1000 * 3600 * 24); + + // Create and register Gauge metrics + Attributes attributes = Attributes.builder() + .put("cert_type", certType) + .put("cert_subject", subject) + .build(); + + ObservableLongGauge observableCertExpireMills = meter.gaugeBuilder(prefix + "cert_expiry_timestamp") + .setDescription("The expiry timestamp of the TLS certificate") + .setUnit("milliseconds") + .ofLongs() + .buildWithCallback(result -> result.record(expiryDate.getTime(), attributes)); + CERT_OBSERVABLE_LONG_GAUGES.add(observableCertExpireMills); + + ObservableLongGauge observableCertExpireDays = meter.gaugeBuilder(prefix + "cert_days_remaining") + .setDescription("The remaining days until the TLS certificate expires") + .setUnit("days") + .ofLongs() + .buildWithCallback(result -> result.record(daysRemaining, attributes)); + CERT_OBSERVABLE_LONG_GAUGES.add(observableCertExpireDays); + } + + /** + * Parse the PEM formatted certificate content into an array of X509 certificates. + * + * @param pemContent The PEM formatted certificate content. + * @return An array of X509 certificates. + * @throws CertificateException If there is an error parsing the certificates. + */ + private static X509Certificate[] parseCertificates(String pemContent) throws CertificateException { + String[] pemArray = pemContent.split("-----END CERTIFICATE-----"); + CertificateFactory factory = CertificateFactory.getInstance("X.509"); + List certList = new ArrayList<>(); + + for (String pemPart : pemArray) { + // Clean the PEM part by removing headers and all whitespace characters + String cleanedPemPart = pemPart.replace("-----BEGIN CERTIFICATE-----", "") + .replaceAll("\\s", ""); // Remove all whitespace characters (spaces, tabs, newlines, etc.) + + // Skip empty parts that might result from splitting + if (cleanedPemPart.isEmpty()) { + continue; + } + + try { + byte[] certBytes = Base64.getDecoder().decode(cleanedPemPart); + X509Certificate cert = (X509Certificate) factory.generateCertificate(new ByteArrayInputStream(certBytes)); + certList.add(cert); + } catch (IllegalArgumentException e) { + LOGGER.warn("Failed to decode certificate part due to invalid Base64, skipping: {}", e.getMessage()); + // Continue processing other certificates instead of failing completely + } catch (CertificateException e) { + LOGGER.warn("Failed to parse certificate, skipping: {}", e.getMessage()); + // Continue processing other certificates instead of failing completely + } + } + + return certList.toArray(new X509Certificate[0]); + } + public static void setIsActiveSupplier(Supplier isActiveSupplier) { S3StreamKafkaMetricsManager.isActiveSupplier = isActiveSupplier; } @@ -258,10 +463,31 @@ public static void setFetchLimiterPermitNumSupplier(Supplier> fetchLimiterWaitingTaskNumSupplier) { + S3StreamKafkaMetricsManager.fetchLimiterWaitingTaskNumSupplier = fetchLimiterWaitingTaskNumSupplier; + } + public static void setFetchPendingTaskNumSupplier(Supplier> fetchPendingTaskNumSupplier) { S3StreamKafkaMetricsManager.fetchPendingTaskNumSupplier = fetchPendingTaskNumSupplier; } + public static CounterMetric buildFetchLimiterTimeoutMetric(String limiterName) { + synchronized (BASE_ATTRIBUTES_LISTENERS) { + CounterMetric metric = new CounterMetric(metricsConfig, FETCH_LIMITER_ATTRIBUTES.get(limiterName), () -> fetchLimiterTimeoutCount); + BASE_ATTRIBUTES_LISTENERS.add(metric); + return metric; + } + } + + public static HistogramMetric buildFetchLimiterTimeMetric(MetricsLevel metricsLevel, String limiterName) { + synchronized (BASE_ATTRIBUTES_LISTENERS) { + HistogramMetric metric = new HistogramMetric(metricsLevel, metricsConfig, FETCH_LIMITER_ATTRIBUTES.get(limiterName)); + BASE_ATTRIBUTES_LISTENERS.add(metric); + FETCH_LIMITER_TIME_METRICS.add(metric); + return metric; + } + } + public static void setLogAppendPermitNumSupplier(Supplier logAppendPermitNumSupplier) { S3StreamKafkaMetricsManager.logAppendPermitNumSupplier = logAppendPermitNumSupplier; } @@ -278,4 +504,16 @@ public static void setPartitionStatusStatisticsSupplier(List partitionSt public static void setTopicPartitionCountMetricsSupplier(Supplier topicPartitionCountSupplier) { S3StreamKafkaMetricsManager.topicPartitionCountSupplier = topicPartitionCountSupplier; } + + public static void setBackPressureStateSupplier(Supplier> backPressureStateSupplier) { + S3StreamKafkaMetricsManager.backPressureStateSupplier = backPressureStateSupplier; + } + + public static void setTruststoreCertsSupplier(Supplier truststoreCertsSupplier) { + S3StreamKafkaMetricsManager.truststoreCertsSupplier = truststoreCertsSupplier; + } + + public static void setCertChainSupplier(Supplier certChainSupplier) { + S3StreamKafkaMetricsManager.certChainSupplier = certChainSupplier; + } } diff --git a/server-common/src/main/java/org/apache/kafka/server/record/ErrorsTolerance.java b/server-common/src/main/java/org/apache/kafka/server/record/ErrorsTolerance.java new file mode 100644 index 0000000000..0740cb79d0 --- /dev/null +++ b/server-common/src/main/java/org/apache/kafka/server/record/ErrorsTolerance.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.server.record; + +import java.util.List; +import java.util.Locale; +import java.util.stream.Collectors; + +import static java.util.Arrays.asList; + +public enum ErrorsTolerance { + NONE("none"), + INVALID_DATA("invalid_data"), + ALL("all"); + + private static final List VALUES = asList(values()); + public final String name; + + ErrorsTolerance(String name) { + this.name = name; + } + + public static List names() { + return VALUES.stream().map(v -> v.name).collect(Collectors.toList()); + } + + public static ErrorsTolerance forName(String name) { + if (name == null) { + return INVALID_DATA; + } + String upperCaseName = name.toUpperCase(Locale.ROOT); + try { + return valueOf(upperCaseName); + } catch (IllegalArgumentException e) { + return INVALID_DATA; + } + } +} diff --git a/server-common/src/main/java/org/apache/kafka/server/record/TableTopicConvertType.java b/server-common/src/main/java/org/apache/kafka/server/record/TableTopicConvertType.java new file mode 100644 index 0000000000..ef8bf17fc4 --- /dev/null +++ b/server-common/src/main/java/org/apache/kafka/server/record/TableTopicConvertType.java @@ -0,0 +1,49 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.server.record; + +import java.util.List; +import java.util.Locale; +import java.util.stream.Collectors; + +import static java.util.Arrays.asList; + +public enum TableTopicConvertType { + RAW("raw"), + STRING("string"), + BY_SCHEMA_ID("by_schema_id"), + BY_LATEST_SCHEMA("by_latest_schema"); + + public final String name; + private static final List VALUES = asList(values()); + + TableTopicConvertType(String name) { + this.name = name; + } + public static List names() { + return VALUES.stream().map(v -> v.name).collect(Collectors.toList()); + } + + public static TableTopicConvertType forName(String n) { + String name = n.toLowerCase(Locale.ROOT); + return VALUES.stream().filter(v -> v.name.equals(name)).findFirst().orElseThrow(() -> + new IllegalArgumentException("Unknown table topic type name: " + name) + ); + } +} diff --git a/server-common/src/main/java/org/apache/kafka/server/record/TableTopicSchemaType.java b/server-common/src/main/java/org/apache/kafka/server/record/TableTopicSchemaType.java new file mode 100644 index 0000000000..abc865e5c3 --- /dev/null +++ b/server-common/src/main/java/org/apache/kafka/server/record/TableTopicSchemaType.java @@ -0,0 +1,50 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.server.record; + +import java.util.List; +import java.util.Locale; +import java.util.stream.Collectors; + +import static java.util.Arrays.asList; + +public enum TableTopicSchemaType { + NONE("none"), + SCHEMALESS("schemaless"), + SCHEMA("schema"); + + public final String name; + private static final List VALUES = asList(values()); + + TableTopicSchemaType(String name) { + this.name = name; + } + + public static List names() { + return VALUES.stream().map(v -> v.name).collect(Collectors.toList()); + } + + public static TableTopicSchemaType forName(String n) { + String name = n.toLowerCase(Locale.ROOT); + return VALUES.stream().filter(v -> v.name.equals(name)).findFirst().orElseThrow(() -> + new IllegalArgumentException("Unknown table topic type name: " + name) + ); + } +} diff --git a/server-common/src/main/java/org/apache/kafka/server/record/TableTopicTransformType.java b/server-common/src/main/java/org/apache/kafka/server/record/TableTopicTransformType.java new file mode 100644 index 0000000000..42997f2b43 --- /dev/null +++ b/server-common/src/main/java/org/apache/kafka/server/record/TableTopicTransformType.java @@ -0,0 +1,49 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.server.record; + +import java.util.List; +import java.util.Locale; +import java.util.stream.Collectors; + +import static java.util.Arrays.asList; + +public enum TableTopicTransformType { + NONE("none"), + FLATTEN("flatten"), + FLATTEN_DEBEZIUM("flatten_debezium"); + + public final String name; + private static final List VALUES = asList(values()); + + TableTopicTransformType(String name) { + this.name = name; + } + public static List names() { + return VALUES.stream().map(v -> v.name).collect(Collectors.toList()); + } + + public static TableTopicTransformType forName(String n) { + String name = n.toLowerCase(Locale.ROOT); + return VALUES.stream().filter(v -> v.name.equals(name)).findFirst().orElseThrow(() -> + new IllegalArgumentException("Unknown table topic type name: " + name) + ); + } +} diff --git a/server-common/src/test/java/org/apache/kafka/server/common/automq/TableTopicConfigValidatorTest.java b/server-common/src/test/java/org/apache/kafka/server/common/automq/TableTopicConfigValidatorTest.java new file mode 100644 index 0000000000..53eeb4f7bf --- /dev/null +++ b/server-common/src/test/java/org/apache/kafka/server/common/automq/TableTopicConfigValidatorTest.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.server.common.automq; + +import org.apache.kafka.common.config.ConfigException; +import org.apache.kafka.server.common.automq.TableTopicConfigValidator.IdColumnsValidator; +import org.apache.kafka.server.common.automq.TableTopicConfigValidator.PartitionValidator; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.apache.kafka.server.common.automq.TableTopicConfigValidator.COMMA_NO_PARENS_REGEX; +import static org.junit.jupiter.api.Assertions.assertThrowsExactly; + +public class TableTopicConfigValidatorTest { + + @Test + public void testPartitionValidator() { + // not an array + assertThrowsExactly(ConfigException.class, () -> PartitionValidator.INSTANCE.ensureValid("config_name", "year(column_name)")); + // invalid transform function + assertThrowsExactly(ConfigException.class, () -> PartitionValidator.INSTANCE.ensureValid("config_name", "[clock(column_name)]")); + // invalid column name + assertThrowsExactly(ConfigException.class, () -> PartitionValidator.INSTANCE.ensureValid("config_name", "[year(\"column_name\")]")); + // invalid bucket + assertThrowsExactly(ConfigException.class, () -> PartitionValidator.INSTANCE.ensureValid("config_name", "[bucket(column_name, abc)]")); + // invalid truncate + assertThrowsExactly(ConfigException.class, () -> PartitionValidator.INSTANCE.ensureValid("config_name", "[truncate(column_name, abc)]")); + + // valid + PartitionValidator.INSTANCE.ensureValid("config_name", "[year(l1.l2.c1), month(c2), day(c3), hour(c4), bucket(c5, 1), truncate(c6, 10)]"); + } + + @Test + public void testIdColumnsValidator() { + // not an array + assertThrowsExactly(ConfigException.class, () -> IdColumnsValidator.INSTANCE.ensureValid("config_name", "c1, c2")); + // invalid column name + assertThrowsExactly(ConfigException.class, () -> IdColumnsValidator.INSTANCE.ensureValid("config_name", "[\"c1\", \"c2\"]")); + + // valid + IdColumnsValidator.INSTANCE.ensureValid("config_name", "[l1.c1, c2]"); + } + + @Test + public void testStringToList() { + Assertions.assertEquals(List.of("a", "b", "c"), TableTopicConfigValidator.stringToList("a, b, c", COMMA_NO_PARENS_REGEX)); + Assertions.assertEquals(List.of("a", "b", "c"), TableTopicConfigValidator.stringToList("[a, b, c]", COMMA_NO_PARENS_REGEX)); + } +} diff --git a/server-common/src/test/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManagerTest.java b/server-common/src/test/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManagerTest.java new file mode 100644 index 0000000000..94d5217d94 --- /dev/null +++ b/server-common/src/test/java/org/apache/kafka/server/metrics/s3stream/S3StreamKafkaMetricsManagerTest.java @@ -0,0 +1,81 @@ +/* + * Copyright 2025, AutoMQ HK Limited. + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.server.metrics.s3stream; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Method; +import java.security.cert.X509Certificate; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +public class S3StreamKafkaMetricsManagerTest { + + @Test + public void testParseCertificatesWithEmptyString() throws Exception { + X509Certificate[] certificates = callParseCertificates(""); + + assertNotNull(certificates); + assertEquals(0, certificates.length); + } + + @Test + public void testParseCertificatesWithWhitespaceInBase64() throws Exception { + // Test certificate with whitespace in Base64 content that would cause "Illegal base64 character 20" error + String certWithSpaces = "-----BEGIN CERTIFICATE-----\n" + + "TUlJQmtUQ0IrUFNKQnFaUUhpUWxDd0ZBTUJReEVqQVFCZ05W" + // base64 line with spaces + " QkFNTUNXeHZZMkZzYUc5emREQWVGdzB5TlRFd01qbHhNREF3TUZG\n" + // Leading space + "QUFNVUNXeHZZMkZzYUc5emREQWVGdzB5TlRFd01qbHhNREF3\t" + // Trailing tab + "TUZGUUFNVUNXeHZZMG\r\n" + // Carriage return + newline + "-----END CERTIFICATE-----"; + + // This should not throw IllegalArgumentException due to the fix + assertDoesNotThrow(() -> { + X509Certificate[] certificates = callParseCertificates(certWithSpaces); + assertNotNull(certificates); + // The certificate might not be valid (just test data), but at least it shouldn't crash with Base64 error + }); + } + + @Test + public void testParseCertificatesWithInvalidBase64() throws Exception { + String invalidCert = "-----BEGIN CERTIFICATE-----\n" + + "InvalidBase64Content!!!\n" + + "-----END CERTIFICATE-----"; + + // Should not throw exception but return empty array due to graceful error handling + assertDoesNotThrow(() -> { + X509Certificate[] certificates = callParseCertificates(invalidCert); + assertNotNull(certificates); + assertEquals(0, certificates.length); // Invalid cert should be skipped + }); + } + + /** + * Helper method to call the private parseCertificates method using reflection + */ + private X509Certificate[] callParseCertificates(String pemContent) throws Exception { + Method method = S3StreamKafkaMetricsManager.class.getDeclaredMethod("parseCertificates", String.class); + method.setAccessible(true); + return (X509Certificate[]) method.invoke(null, pemContent); + } +} diff --git a/server/src/main/java/org/apache/kafka/network/SocketServerConfigs.java b/server/src/main/java/org/apache/kafka/network/SocketServerConfigs.java index 72c4bb53d1..b4696a9f4f 100644 --- a/server/src/main/java/org/apache/kafka/network/SocketServerConfigs.java +++ b/server/src/main/java/org/apache/kafka/network/SocketServerConfigs.java @@ -158,6 +158,12 @@ public class SocketServerConfigs { public static final int QUEUED_MAX_REQUESTS_DEFAULT = 500; public static final String QUEUED_MAX_REQUESTS_DOC = "The number of queued requests allowed for data-plane, before blocking the network threads"; + // AutoMQ inject start + public static final String QUEUED_MAX_REQUESTS_SIZE_BYTES_CONFIG = "queued.max.requests.size.bytes"; + public static final int QUEUED_MAX_REQUESTS_SIZE_BYTES_DEFAULT = 100 * 1024 * 1024; + public static final String QUEUED_MAX_REQUESTS_SIZE_BYTES_DOC = "The number of queued requests size in total allowed for data-plane, before blocking the network threads"; + // AutoMQ inject end + public static final String QUEUED_MAX_BYTES_CONFIG = "queued.max.request.bytes"; public static final int QUEUED_MAX_REQUEST_BYTES_DEFAULT = -1; public static final String QUEUED_MAX_REQUEST_BYTES_DOC = "The number of queued bytes allowed before no more requests are read"; @@ -182,6 +188,9 @@ public class SocketServerConfigs { .define(CONNECTIONS_MAX_IDLE_MS_CONFIG, LONG, CONNECTIONS_MAX_IDLE_MS_DEFAULT, MEDIUM, CONNECTIONS_MAX_IDLE_MS_DOC) .define(FAILED_AUTHENTICATION_DELAY_MS_CONFIG, INT, FAILED_AUTHENTICATION_DELAY_MS_DEFAULT, atLeast(0), LOW, FAILED_AUTHENTICATION_DELAY_MS_DOC) .define(QUEUED_MAX_REQUESTS_CONFIG, INT, QUEUED_MAX_REQUESTS_DEFAULT, atLeast(1), HIGH, QUEUED_MAX_REQUESTS_DOC) + // AutoMQ inject start + .define(QUEUED_MAX_REQUESTS_SIZE_BYTES_CONFIG, INT, QUEUED_MAX_REQUESTS_SIZE_BYTES_DEFAULT, atLeast(1024 * 1024), HIGH, QUEUED_MAX_REQUESTS_SIZE_BYTES_DOC) + // AutoMQ inject end .define(QUEUED_MAX_BYTES_CONFIG, LONG, QUEUED_MAX_REQUEST_BYTES_DEFAULT, MEDIUM, QUEUED_MAX_REQUEST_BYTES_DOC) .define(NUM_NETWORK_THREADS_CONFIG, INT, NUM_NETWORK_THREADS_DEFAULT, atLeast(1), HIGH, NUM_NETWORK_THREADS_DOC); } diff --git a/server/src/main/java/org/apache/kafka/server/config/BrokerQuotaManagerConfig.java b/server/src/main/java/org/apache/kafka/server/config/BrokerQuotaManagerConfig.java index 53538e55c7..518d23aa6a 100644 --- a/server/src/main/java/org/apache/kafka/server/config/BrokerQuotaManagerConfig.java +++ b/server/src/main/java/org/apache/kafka/server/config/BrokerQuotaManagerConfig.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.server.config; @@ -26,7 +34,8 @@ public class BrokerQuotaManagerConfig extends ClientQuotaManagerConfig { private boolean quotaEnabled = false; private double produceQuota = Double.MAX_VALUE; private double fetchQuota = Double.MAX_VALUE; - private double requestQuota = Double.MAX_VALUE; + private double slowFetchQuota = Double.MAX_VALUE; + private double requestRateQuota = Double.MAX_VALUE; private List userWhiteList = List.of(); private List clientIdWhiteList = List.of(); @@ -42,7 +51,8 @@ public void update(Properties props) { quotaEnabled = getBoolean(map, QuotaConfigs.BROKER_QUOTA_ENABLED_CONFIG, quotaEnabled); produceQuota = getDouble(map, QuotaConfigs.BROKER_QUOTA_PRODUCE_BYTES_CONFIG, produceQuota); fetchQuota = getDouble(map, QuotaConfigs.BROKER_QUOTA_FETCH_BYTES_CONFIG, fetchQuota); - requestQuota = getDouble(map, QuotaConfigs.BROKER_QUOTA_REQUEST_RATE_CONFIG, requestQuota); + slowFetchQuota = getDouble(map, QuotaConfigs.BROKER_QUOTA_SLOW_FETCH_BYTES_CONFIG, slowFetchQuota); + requestRateQuota = getDouble(map, QuotaConfigs.BROKER_QUOTA_REQUEST_RATE_CONFIG, requestRateQuota); String userWhiteListProp = props.getProperty(QuotaConfigs.BROKER_QUOTA_WHITE_LIST_USER_CONFIG); if (null != userWhiteListProp && !userWhiteListProp.isBlank()) { @@ -72,12 +82,32 @@ public double produceQuota() { return produceQuota; } + public void produceQuota(double produceQuota) { + this.produceQuota = produceQuota; + } + public double fetchQuota() { return fetchQuota; } - public double requestQuota() { - return requestQuota; + public void fetchQuota(double fetchQuota) { + this.fetchQuota = fetchQuota; + } + + public double slowFetchQuota() { + return slowFetchQuota; + } + + public void slowFetchQuota(double slowFetchQuota) { + this.slowFetchQuota = slowFetchQuota; + } + + public double requestRateQuota() { + return requestRateQuota; + } + + public void requestRateQuota(double requestRateQuota) { + this.requestRateQuota = requestRateQuota; } public List userWhiteList() { diff --git a/settings.gradle b/settings.gradle index 74582b5937..99554bfe89 100644 --- a/settings.gradle +++ b/settings.gradle @@ -103,7 +103,9 @@ include 'clients', 'transaction-coordinator', 'trogdor', 's3stream', - 'automq-shell' + 'automq-shell', + 'automq-log-uploader', + 'automq-metrics' project(":storage:api").name = "storage-api" rootProject.name = 'kafka' diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogConfig.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogConfig.java index a2c37f081b..fa28e49ce9 100644 --- a/storage/src/main/java/org/apache/kafka/storage/internals/log/LogConfig.java +++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/LogConfig.java @@ -32,14 +32,20 @@ import org.apache.kafka.common.record.RecordVersion; import org.apache.kafka.common.record.Records; import org.apache.kafka.common.record.TimestampType; +import org.apache.kafka.common.requests.ListOffsetsRequest; import org.apache.kafka.common.utils.ConfigUtils; import org.apache.kafka.common.utils.Utils; import org.apache.kafka.server.common.MetadataVersion; import org.apache.kafka.server.common.MetadataVersionValidator; +import org.apache.kafka.server.common.automq.TableTopicConfigValidator; import org.apache.kafka.server.config.QuotaConfigs; import org.apache.kafka.server.config.ServerLogConfigs; import org.apache.kafka.server.config.ServerTopicConfigSynonyms; import org.apache.kafka.server.record.BrokerCompressionType; +import org.apache.kafka.server.record.ErrorsTolerance; +import org.apache.kafka.server.record.TableTopicConvertType; +import org.apache.kafka.server.record.TableTopicSchemaType; +import org.apache.kafka.server.record.TableTopicTransformType; import java.util.Collections; import java.util.HashMap; @@ -68,6 +74,11 @@ import static org.apache.kafka.common.config.ConfigDef.Type.STRING; import static org.apache.kafka.common.config.ConfigDef.ValidString.in; import static org.apache.kafka.server.common.MetadataVersion.IBP_3_0_IV1; +import static org.apache.kafka.server.record.TableTopicConvertType.BY_LATEST_SCHEMA; +import static org.apache.kafka.server.record.TableTopicConvertType.BY_SCHEMA_ID; +import static org.apache.kafka.server.record.TableTopicConvertType.RAW; +import static org.apache.kafka.server.record.TableTopicTransformType.FLATTEN_DEBEZIUM; +import static org.apache.kafka.server.record.TableTopicTransformType.NONE; public class LogConfig extends AbstractConfig { @@ -253,6 +264,11 @@ public Optional serverConfigName(String configName) { .define(ServerLogConfigs.ALTER_CONFIG_POLICY_CLASS_NAME_CONFIG, CLASS, null, LOW, ServerLogConfigs.ALTER_CONFIG_POLICY_CLASS_NAME_DOC) .define(ServerLogConfigs.LOG_MESSAGE_DOWNCONVERSION_ENABLE_CONFIG, BOOLEAN, ServerLogConfigs.LOG_MESSAGE_DOWNCONVERSION_ENABLE_DEFAULT, LOW, ServerLogConfigs.LOG_MESSAGE_DOWNCONVERSION_ENABLE_DOC) .define(ServerLogConfigs.LOG_DIR_FAILURE_TIMEOUT_MS_CONFIG, LONG, ServerLogConfigs.LOG_DIR_FAILURE_TIMEOUT_MS_DEFAULT, atLeast(1), LOW, ServerLogConfigs.LOG_DIR_FAILURE_TIMEOUT_MS_DOC) + + // AutoMQ inject start + .define(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG, STRING, null, null, MEDIUM, TopicConfig.TABLE_TOPIC_NAMESPACE_DOC) + // AutoMQ inject end + .defineInternal(ServerLogConfigs.LOG_INITIAL_TASK_DELAY_MS_CONFIG, LONG, ServerLogConfigs.LOG_INITIAL_TASK_DELAY_MS_DEFAULT, atLeast(0), LOW, ServerLogConfigs.LOG_INITIAL_TASK_DELAY_MS_DOC); private static final LogConfigDef CONFIG = new LogConfigDef(); @@ -328,6 +344,28 @@ public Optional serverConfigName(String configName) { // AutoMQ inject start // dynamic config #validateNames check will check the old configs name validity .define("elasticstream.replication.factor", INT, 1, atLeast(1), LOW, "deprecated, should not remove for compatibility") + .define(TopicConfig.TABLE_TOPIC_ENABLE_CONFIG, BOOLEAN, false, null, MEDIUM, TopicConfig.TABLE_TOPIC_ENABLE_DOC) + .define(TopicConfig.TABLE_TOPIC_COMMIT_INTERVAL_CONFIG, LONG, TimeUnit.MINUTES.toMillis(1), between(1, TimeUnit.MINUTES.toMillis(15)), MEDIUM, TopicConfig.TABLE_TOPIC_COMMIT_INTERVAL_DOC) + .define(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG, STRING, null, null, MEDIUM, TopicConfig.TABLE_TOPIC_NAMESPACE_DOC) + .define(TopicConfig.TABLE_TOPIC_SCHEMA_TYPE_CONFIG, STRING, TableTopicSchemaType.NONE.name, in(TableTopicSchemaType.names().toArray(new String[0])), MEDIUM, TopicConfig.TABLE_TOPIC_SCHEMA_TYPE_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG, STRING, RAW.name, in(TableTopicConvertType.names().toArray(new String[0])), MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG, STRING, TableTopicConvertType.STRING.name, in(TableTopicConvertType.names().toArray(new String[0])), MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG, STRING, NONE.name, in(TableTopicTransformType.names().toArray(new String[0])), MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_SUBJECT_CONFIG, STRING, null, null, MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_SUBJECT_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_CONFIG, STRING, null, null, MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_SUBJECT_CONFIG, STRING, null, null, MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_SUBJECT_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_CONFIG, STRING, null, null, MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_BY_LATEST_SCHEMA_MESSAGE_FULL_NAME_DOC) + .define(TopicConfig.TABLE_TOPIC_ID_COLUMNS_CONFIG, STRING, null, TableTopicConfigValidator.IdColumnsValidator.INSTANCE, MEDIUM, TopicConfig.TABLE_TOPIC_ID_COLUMNS_DOC) + .define(TopicConfig.TABLE_TOPIC_PARTITION_BY_CONFIG, STRING, null, TableTopicConfigValidator.PartitionValidator.INSTANCE, MEDIUM, TopicConfig.TABLE_TOPIC_PARTITION_BY_DOC) + .define(TopicConfig.TABLE_TOPIC_UPSERT_ENABLE_CONFIG, BOOLEAN, false, null, MEDIUM, TopicConfig.TABLE_TOPIC_UPSERT_ENABLE_DOC) + .define(TopicConfig.TABLE_TOPIC_CDC_FIELD_CONFIG, STRING, null, null, MEDIUM, TopicConfig.TABLE_TOPIC_CDC_FIELD_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_ERRORS_TOLERANCE_CONFIG, STRING, ErrorsTolerance.INVALID_DATA.name, in(ErrorsTolerance.names().toArray(new String[0])), MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_ERRORS_TOLERANCE_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_CONFIG, BOOLEAN, TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_DEFAULT, MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_CONFIG, INT, TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_DEFAULT, atLeast(1), MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_DOC) + .define(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_CONFIG, INT, TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_DEFAULT, atLeast(1), MEDIUM, TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_DOC) + .define(TopicConfig.KAFKA_LINKS_ID_CONFIG, STRING, null, null, MEDIUM, TopicConfig.KAFKA_LINKS_ID_DOC) + .define(TopicConfig.KAFKA_LINKS_TOPIC_START_TIME_CONFIG, LONG, ListOffsetsRequest.LATEST_TIMESTAMP, null, MEDIUM, TopicConfig.KAFKA_LINKS_TOPIC_START_TIME_DOC) + .define(TopicConfig.KAFKA_LINKS_TOPIC_STATE_CONFIG, STRING, null, null, MEDIUM, TopicConfig.KAFKA_LINKS_TOPIC_STATE_DOC) // AutoMQ inject end .define(TopicConfig.REMOTE_LOG_DISABLE_POLICY_CONFIG, STRING, TopicConfig.REMOTE_LOG_DISABLE_POLICY_RETAIN, in(TopicConfig.REMOTE_LOG_DISABLE_POLICY_RETAIN, TopicConfig.REMOTE_LOG_DISABLE_POLICY_DELETE), @@ -378,6 +416,31 @@ public Optional serverConfigName(String configName) { public final boolean messageDownConversionEnable; public final RemoteLogConfig remoteLogConfig; + // AutoMQ inject start + public final boolean tableTopicEnable; + public final long tableTopicCommitInterval; + public final String tableTopicNamespace; + @Deprecated + public final TableTopicSchemaType tableTopicSchemaType; + public final TableTopicConvertType valueConvertType; + public final TableTopicConvertType keyConvertType; + + public final TableTopicTransformType transformType; + public final String tableTopicIdColumns; + public final String tableTopicPartitionBy; + public final boolean tableTopicUpsertEnable; + public final String tableTopicCdcField; + public final ErrorsTolerance errorsTolerance; + + public final String kafkaLinksId; + public final long kafkaLinksTopicStartTime; + public final String kafkaLinksTopicState; + + public final boolean tableTopicExpireSnapshotEnabled; + public final int tableTopicExpireSnapshotOlderThanHours; + public final int tableTopicExpireSnapshotRetainLast; + // AutoMQ inject end + private final int maxMessageSize; private final Map props; @@ -428,6 +491,29 @@ public LogConfig(Map props, Set overriddenConfigs) { this.followerReplicationThrottledReplicas = Collections.unmodifiableList(getList(QuotaConfigs.FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG)); this.messageDownConversionEnable = getBoolean(TopicConfig.MESSAGE_DOWNCONVERSION_ENABLE_CONFIG); + // AutoMQ inject start + this.tableTopicEnable = getBoolean(TopicConfig.TABLE_TOPIC_ENABLE_CONFIG); + this.tableTopicCommitInterval = getLong(TopicConfig.TABLE_TOPIC_COMMIT_INTERVAL_CONFIG); + this.tableTopicNamespace = getString(TopicConfig.TABLE_TOPIC_NAMESPACE_CONFIG); + this.tableTopicSchemaType = TableTopicSchemaType.forName(getString(TopicConfig.TABLE_TOPIC_SCHEMA_TYPE_CONFIG)); + this.valueConvertType = TableTopicConvertType.forName(getString(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG)); + this.keyConvertType = TableTopicConvertType.forName(getString(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG)); + this.transformType = TableTopicTransformType.forName(getString(TopicConfig.AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG)); + this.tableTopicIdColumns = getString(TopicConfig.TABLE_TOPIC_ID_COLUMNS_CONFIG); + this.tableTopicPartitionBy = getString(TopicConfig.TABLE_TOPIC_PARTITION_BY_CONFIG); + this.tableTopicUpsertEnable = getBoolean(TopicConfig.TABLE_TOPIC_UPSERT_ENABLE_CONFIG); + this.tableTopicCdcField = getString(TopicConfig.TABLE_TOPIC_CDC_FIELD_CONFIG); + this.errorsTolerance = ErrorsTolerance.forName(getString(TopicConfig.AUTOMQ_TABLE_TOPIC_ERRORS_TOLERANCE_CONFIG)); + + this.kafkaLinksId = getString(TopicConfig.KAFKA_LINKS_ID_CONFIG); + this.kafkaLinksTopicStartTime = getLong(TopicConfig.KAFKA_LINKS_TOPIC_START_TIME_CONFIG); + this.kafkaLinksTopicState = getString(TopicConfig.KAFKA_LINKS_TOPIC_STATE_CONFIG); + + this.tableTopicExpireSnapshotEnabled = getBoolean(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_ENABLED_CONFIG); + this.tableTopicExpireSnapshotOlderThanHours = getInt(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_OLDER_THAN_HOURS_CONFIG); + this.tableTopicExpireSnapshotRetainLast = getInt(TopicConfig.AUTOMQ_TABLE_TOPIC_EXPIRE_SNAPSHOT_RETAIN_LAST_CONFIG); + // AutoMQ inject end + remoteLogConfig = new RemoteLogConfig(this); } @@ -703,6 +789,26 @@ public static void validate(Properties props, } } + // AutoMQ inject start + public static void validateTableTopicSchemaConfigValues(Properties props, String tableTopicSchemaRegistryUrl) { + String valueConvertProperty = props.getProperty(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_VALUE_TYPE_CONFIG); + String keyConvertProperty = props.getProperty(TopicConfig.AUTOMQ_TABLE_TOPIC_CONVERT_KEY_TYPE_CONFIG); + String transformProperty = props.getProperty(TopicConfig.AUTOMQ_TABLE_TOPIC_TRANSFORM_VALUE_TYPE_CONFIG); + + // Validation logic using new (or mapped) configs + if ((BY_SCHEMA_ID.name.equals(valueConvertProperty) || BY_LATEST_SCHEMA.name.equals(valueConvertProperty)) && (tableTopicSchemaRegistryUrl == null || tableTopicSchemaRegistryUrl.isEmpty())) { + throw new InvalidConfigurationException("Table topic convert type is set to '" + valueConvertProperty + "' but schema registry URL is not configured"); + } + if ((BY_SCHEMA_ID.name.equals(keyConvertProperty) || BY_LATEST_SCHEMA.name.equals(keyConvertProperty)) && (tableTopicSchemaRegistryUrl == null || tableTopicSchemaRegistryUrl.isEmpty())) { + throw new InvalidConfigurationException("Table topic convert type is set to '" + keyConvertProperty + "' but schema registry URL is not configured"); + } + if (!(BY_SCHEMA_ID.name.equals(valueConvertProperty) || BY_LATEST_SCHEMA.name.equals(valueConvertProperty)) && FLATTEN_DEBEZIUM.name.equals(transformProperty)) { + throw new InvalidConfigurationException(valueConvertProperty + " convert type cannot be used with '" + FLATTEN_DEBEZIUM.name + "' transform type"); + } + } + + // AutoMQ inject end + @Override public String toString() { return "LogConfig{" + diff --git a/storage/src/main/java/org/apache/kafka/storage/internals/log/TimeIndex.java b/storage/src/main/java/org/apache/kafka/storage/internals/log/TimeIndex.java index 54f0176c8b..044608411d 100644 --- a/storage/src/main/java/org/apache/kafka/storage/internals/log/TimeIndex.java +++ b/storage/src/main/java/org/apache/kafka/storage/internals/log/TimeIndex.java @@ -73,8 +73,10 @@ public TimeIndex(File file, long baseOffset, int maxIndexSize, boolean writable, // AutoMQ inject start if (!noopFile) { this.lastEntry = lastEntryFromIndexFile(); - log.debug("Loaded index file {} with maxEntries = {}, maxIndexSize = {}, entries = {}, lastOffset = {}, file position = {}", - file.getAbsolutePath(), maxEntries(), maxIndexSize, entries(), lastEntry.offset, mmap().position()); + if (log.isDebugEnabled()) { + log.debug("Loaded index file {} with maxEntries = {}, maxIndexSize = {}, entries = {}, lastOffset = {}, file position = {}", + file.getAbsolutePath(), maxEntries(), maxIndexSize, entries(), lastEntry.offset, mmap().position()); + } } // AutoMQ inject end } diff --git a/streams/src/test/java/org/apache/kafka/streams/tests/RelationalSmokeTest.java b/streams/src/test/java/org/apache/kafka/streams/tests/RelationalSmokeTest.java index 1b24fd6175..0569acba9b 100644 --- a/streams/src/test/java/org/apache/kafka/streams/tests/RelationalSmokeTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/tests/RelationalSmokeTest.java @@ -969,8 +969,8 @@ public static void main(final String[] args) { switch (mode) { case "driver": { // this starts the driver (data generation and result verification) - final int numArticles = 1_000; - final int numComments = 10_000; + final int numArticles = 1_00; + final int numComments = 10_0; final DataSet dataSet = DataSet.generate(numArticles, numComments); // publish the data for at least one minute dataSet.produce(kafka, Duration.ofMinutes(1)); diff --git a/tests/docker/Dockerfile b/tests/docker/Dockerfile index e746f79990..a5822a61ee 100644 --- a/tests/docker/Dockerfile +++ b/tests/docker/Dockerfile @@ -71,6 +71,7 @@ RUN pip3 install --upgrade cffi virtualenv pyasn1 boto3 pycrypto pywinrm ipaddre # AutoMQ inject start RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-$(uname -m | sed 's/arm64/aarch64/')-2.15.15.zip" -o "awscliv2.zip" && unzip awscliv2.zip && ./aws/install && rm -rf awscliv2.zip aws +RUN curl -fsSL https://get.docker.com -o get-docker.sh && sudo sh get-docker.sh # AutoMQ inject end COPY --from=build-native-image /build/kafka-binary/ /opt/kafka-binary/ diff --git a/tests/docker/ducker-ak b/tests/docker/ducker-ak index 29975617d3..af1224c7e3 100755 --- a/tests/docker/ducker-ak +++ b/tests/docker/ducker-ak @@ -293,14 +293,23 @@ docker_run() { if [[ -n ${port_mapping} ]]; then expose_ports="${expose_ports} -p ${port_mapping}:${port_mapping}" fi - - # Invoke docker-run. We need privileged mode to be able to run iptables - # and mount FUSE filesystems inside the container. We also need it to - # run iptables inside the container. - must_do -v docker run --privileged \ - -d -t -h "${node}" --network ducknet "${expose_ports}" \ - --memory=${docker_run_memory_limit} --memory-swappiness=1 \ - -v "${kafka_dir}:/opt/kafka-dev" --name "${node}" -- "${image_name}" +# AutoMQ inject start + if [[ "${node}" == "ducker01" ]]; then + must_do -v docker run --privileged \ + -d -t -h "${node}" --network ducknet "${expose_ports}" \ + --memory=${docker_run_memory_limit} --memory-swappiness=1 \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -v "${kafka_dir}:/opt/kafka-dev" --name "${node}" -- "${image_name}" + else +# AutoMQ inject end + # Invoke docker-run. We need privileged mode to be able to run iptables + # and mount FUSE filesystems inside the container. We also need it to + # run iptables inside the container. + must_do -v docker run --privileged \ + -d -t -h "${node}" --network ducknet "${expose_ports}" \ + --memory=${docker_run_memory_limit} --memory-swappiness=1 \ + -v "${kafka_dir}:/opt/kafka-dev" --name "${node}" -- "${image_name}" + fi } setup_custom_ducktape() { diff --git a/tests/docker/iceberg-catalog/docker-compose.yaml b/tests/docker/iceberg-catalog/docker-compose.yaml new file mode 100644 index 0000000000..62e5e24388 --- /dev/null +++ b/tests/docker/iceberg-catalog/docker-compose.yaml @@ -0,0 +1,43 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +services: + rest: + image: apache/iceberg-rest-fixture + container_name: "${REST_DOCKER_NAME-rest}" + hostname: "${REST_DOCKER_NAME-rest}" + ports: + - 8181:8181 + environment: + - AWS_ACCESS_KEY_ID=admin + - AWS_SECRET_ACCESS_KEY=password + - AWS_REGION=us-east-1 + - CATALOG_WAREHOUSE=s3://ko3/iceberg/ + - CATALOG_IO__IMPL=org.apache.iceberg.aws.s3.S3FileIO + - CATALOG_S3_ENDPOINT=http://10.5.0.2:4566 + healthcheck: + test: ["CMD-SHELL", "curl -f http://localhost:8181/v1/config || exit 1"] + interval: 10s + timeout: 5s + retries: 20 + start_period: 30s + networks: + ducknet: + ipv4_address: 10.5.1.4 + +networks: + ducknet: + name: ducknet + external: true diff --git a/tests/docker/schema-registry/docker-compose.yaml b/tests/docker/schema-registry/docker-compose.yaml new file mode 100644 index 0000000000..a62b07ef1b --- /dev/null +++ b/tests/docker/schema-registry/docker-compose.yaml @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +services: + schema-registry: + image: confluentinc/cp-schema-registry:latest + container_name: "${SCHEMA_REGISTRY_DOCKER_NAME-schema-registry}" + hostname: "${SCHEMA_REGISTRY_DOCKER_NAME-schema-registry}" + ports: + - "8081:8081" + environment: + SCHEMA_REGISTRY_HOST_NAME: ${LOCALSTACK_DOCKER_NAME-schema-registry} + SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS: ${KAFKA_BOOTSTRAP_SERVERS} + SCHEMA_REGISTRY_LISTENERS: http://0.0.0.0:8081 + restart: on-failure + healthcheck: + test: ["CMD-SHELL", "curl -f http://localhost:8081 || exit 1"] + interval: 10s + timeout: 60s + retries: 30 + networks: + ducknet: + ipv4_address: 10.5.1.3 + +networks: + ducknet: + name: ducknet + external: true \ No newline at end of file diff --git a/tests/kafkatest/automq/__init__.py b/tests/kafkatest/automq/__init__.py index ba3900e295..dad893e0d4 100644 --- a/tests/kafkatest/automq/__init__.py +++ b/tests/kafkatest/automq/__init__.py @@ -1,11 +1,4 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ diff --git a/tests/kafkatest/automq/autobalancer_test.py b/tests/kafkatest/automq/autobalancer_test.py deleted file mode 100644 index c0a0e604f2..0000000000 --- a/tests/kafkatest/automq/autobalancer_test.py +++ /dev/null @@ -1,255 +0,0 @@ -""" -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 -""" - -from ducktape.mark import parametrize -from ducktape.mark.resource import cluster -from ducktape.tests.test import Test -from ducktape.mark import matrix -from kafkatest.automq.automq_e2e_util import (FILE_WAL, S3_WAL, run_simple_load, TOPIC, append_info) -from kafkatest.services.kafka import KafkaService - -# Configuration constants for the AutoBalancer -REPORT_INTERVAL = 'autobalancer.reporter.metrics.reporting.interval.ms' -DETECT_INTERVAL = 'autobalancer.controller.anomaly.detect.interval.ms' -ENABLE = 'autobalancer.controller.enable' -IN_AVG_DEVIATION = 'autobalancer.controller.network.in.distribution.detect.avg.deviation' -OUT_AVG_DEVIATION = 'autobalancer.controller.network.out.distribution.detect.avg.deviation' -OUT_THRESHOLD = 'autobalancer.controller.network.out.usage.distribution.detect.threshold' -IN_THRESHOLD = 'autobalancer.controller.network.in.usage.distribution.detect.threshold' -GOALS = 'autobalancer.controller.goals' -EXCLUDE_TOPIC = 'autobalancer.controller.exclude.topics' -EXCLUDE_BROKER = 'autobalancer.controller.exclude.broker.ids' -METRIC_REPORTERS = 'metric.reporters' - - -def get_partition_count_per_broker(partition_data, exclude=None): - """ - Get the count of partitions per broker. - - :param partition_data: Data of partitions - :param exclude: Set of broker IDs to exclude from the count - :return: Dictionary with broker IDs as keys and partition counts as values - """ - if exclude is None: - exclude = set() - broker_replicas_count = {} - - partitions = partition_data.get('partitions', []) - for partition in partitions: - replicas = partition.get('replicas', []) - for broker_id in replicas: - broker_id = str(broker_id) - if broker_id in exclude: - continue - if broker_id not in broker_replicas_count: - broker_replicas_count[broker_id] = 0 - broker_replicas_count[broker_id] += 1 - - return broker_replicas_count - - -def check_partition_replicas(partition_data, exclude=None): - """ - Check if the partition replicas are evenly distributed among brokers. - - :param partition_data: Data of partitions - :param exclude: Set of broker IDs to exclude from the check - :return: Tuple of success status and message - """ - broker_replicas_count = get_partition_count_per_broker(partition_data, exclude) - replica_counts = list(broker_replicas_count.values()) - - success = True - msg = '' - if len(set(replica_counts)) != 1: - success = False - msg = "Brokers have different numbers of partition replicas: " + str(broker_replicas_count) - - return success, msg - - -def check_partition_eq(topic_info1, topic_info2): - success, msg = True, '' - success = str(topic_info1) == str(topic_info2) - topic_name = topic_info1.get('partitions')[0]['topic'] - if not success: - msg += f"Topic {topic_name} was modified despite being excluded from AutoBalancer. Before: {topic_info1}, After: {topic_info2}" - return success, msg - - -class AutoBalancerTest(Test): - """ - Test class for AutoBalancer functionality - """ - - def __init__(self, test_context): - super(AutoBalancerTest, self).__init__(test_context) - self.context = test_context - self.start = False - self.topic = TOPIC - self.avg_deviation = 0.2 - self.maximum_broker_deviation_percentage = 0.15 - - def create_kafka(self, num_nodes=1, partition=1, exclude_broker=None, exclude_topic=None, replica_assignment=None, wal='file'): - """ - Create and configure a Kafka cluster for testing. - - :param num_nodes: Number of Kafka nodes - :param partition: Number of partitions - :param exclude_broker: Brokers to exclude from AutoBalancer - :param exclude_topic: Topics to exclude from AutoBalancer - :param replica_assignment: Replica assignment for partitions - """ - log_size = 256 * 1024 * 1024 - block_size = 256 * 1024 * 1024 - threshold = 512 * 1024 - server_prop_overrides = [ - ['s3.wal.cache.size', str(log_size)], - ['s3.wal.capacity', str(log_size)], - ['s3.wal.upload.threshold', str(log_size // 4)], - ['s3.block.cache.size', str(block_size)], - [ENABLE, 'true'], - [IN_AVG_DEVIATION, str(self.avg_deviation)], - [OUT_AVG_DEVIATION, str(self.avg_deviation)], - [GOALS, - 'kafka.autobalancer.goals.NetworkInUsageDistributionGoal,' - 'kafka.autobalancer.goals.NetworkOutUsageDistributionGoal'], - [IN_THRESHOLD, str(threshold)], - [OUT_THRESHOLD, str(threshold)], - [REPORT_INTERVAL, str(4000)], - [DETECT_INTERVAL, str(8000)], - [METRIC_REPORTERS, 'kafka.autobalancer.metricsreporter.AutoBalancerMetricsReporter'], - ['s3.wal.path', FILE_WAL if wal == 'file' else S3_WAL], - ] - - if exclude_broker: - server_prop_overrides.append([EXCLUDE_BROKER, exclude_broker]) - - if exclude_topic: - server_prop_overrides.append([EXCLUDE_TOPIC, exclude_topic]) - - self.controller_num_nodes_override = 0 - if num_nodes == 3: - self.controller_num_nodes_override = 1 # only use one combined node - - self.kafka = KafkaService(self.context, num_nodes=num_nodes, zk=None, - kafka_heap_opts="-Xmx2048m -Xms2048m", - server_prop_overrides=server_prop_overrides, - topics={ - self.topic: { - 'partitions': partition, - 'replication-factor': 1, - "replica-assignment": replica_assignment, - 'configs': { - 'min.insync.replicas': 1, - } - }, - }, - controller_num_nodes_override=self.controller_num_nodes_override - ) - self.start = True - - @cluster(num_nodes=5) - @matrix(automq_num_nodes=[2], partition=[4], replica_assignment=['1,1,1,2'], wal=['file', 's3']) - def test_action(self, automq_num_nodes, partition, replica_assignment, wal): - """ - Test throughput distribution across brokers - :param automq_num_nodes: Number of automq - :param partition: Number of partitions - :param replica_assignment: Replica assignment for partitions - """ - success, msg = True, '' - self.create_kafka(num_nodes=automq_num_nodes, partition=partition, replica_assignment=replica_assignment, wal=wal) - self.kafka.start() - - run_simple_load(test_context=self.context, kafka=self.kafka, logger=self.logger, topic=self.topic, - num_records=20000, throughput=1300) - - topic_after = self.kafka.parse_describe_topic(self.kafka.describe_topic(TOPIC)) - success_, msg_ = check_partition_replicas(topic_after) - success = success and success_ - msg = append_info(msg, success_, msg_) - - assert success, msg - - @cluster(num_nodes=4) - @matrix(automq_num_nodes=[2], exclude_broker=['2'], partition=[4], replica_assignment=['1,1,1,2'], wal=['file', 's3']) - def test_broker_white_list(self, automq_num_nodes, exclude_broker, partition, replica_assignment, wal): - """ - Test broker exclusion functionality - :param automq_num_nodes: Number of automq - :param exclude_broker: Brokers to exclude from AutoBalancer - :param partition: Number of partitions - :param replica_assignment: Replica assignment for partitions - """ - success, msg = True, '' - self.create_kafka(num_nodes=automq_num_nodes, exclude_broker=exclude_broker, partition=partition, - replica_assignment=replica_assignment, wal=wal) - self.kafka.start() - before = self.kafka.parse_describe_topic(self.kafka.describe_topic(TOPIC)) - run_simple_load(test_context=self.context, kafka=self.kafka, logger=self.logger, topic=self.topic, - num_records=20000, throughput=1300) - after = self.kafka.parse_describe_topic(self.kafka.describe_topic(TOPIC)) - - success_, msg_ = check_partition_eq(topic_info1=before, - topic_info2=after) - success = success and success_ - msg = append_info(msg, success_, msg_) - - assert success, msg - - @cluster(num_nodes=6) - @matrix(automq_num_nodes=[2], wal=['file', 's3']) - def test_topic_white_list(self, automq_num_nodes, wal): - """ - Test topic exclusion functionality - :param automq_num_nodes: Number of automq - """ - success, msg = True, '' - topic1 = 'test_topic01' - topic_cfg1 = { - "topic": topic1, - "partitions": 4, - "replication-factor": 1, - "configs": {"min.insync.replicas": 1}, - "replica-assignment": '1,1,1,2', - } - topic2 = 'test_topic02' - topic_cfg2 = { - "topic": topic2, - "partitions": 4, - "replication-factor": 1, - "configs": {"min.insync.replicas": 1}, - "replica-assignment": '1,1,1,2', - } - self.create_kafka(num_nodes=automq_num_nodes, exclude_topic=topic1, partition=1, replica_assignment='1', wal=wal) - self.kafka.start() - self.kafka.create_topic(topic_cfg1) - self.kafka.create_topic(topic_cfg2) - - topic1_before = self.kafka.parse_describe_topic(self.kafka.describe_topic(topic1)) - run_simple_load(test_context=self.context, kafka=self.kafka, logger=self.logger, topic=topic1, - num_records=15000, throughput=1300) - run_simple_load(test_context=self.context, kafka=self.kafka, logger=self.logger, topic=topic2, - num_records=15000, throughput=1300) - - topic1_after = self.kafka.parse_describe_topic(self.kafka.describe_topic(topic1)) - topic2_after = self.kafka.parse_describe_topic(self.kafka.describe_topic(topic2)) - - success_, msg_ = check_partition_eq(topic1_before, topic1_after) - success = success and success_ - msg = append_info(msg, success_, msg_) - - success_, msg_ = check_partition_replicas(topic2_after) - success = success and success_ - msg = append_info(msg, success_, msg_) - - assert success, msg diff --git a/tests/kafkatest/automq/automq_e2e_util.py b/tests/kafkatest/automq/automq_e2e_util.py index 98ecb2bb6f..60d372fe87 100644 --- a/tests/kafkatest/automq/automq_e2e_util.py +++ b/tests/kafkatest/automq/automq_e2e_util.py @@ -1,12 +1,5 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ import re diff --git a/tests/kafkatest/automq/compaction_test.py b/tests/kafkatest/automq/compaction_test.py index 52dccce70f..a6702204be 100644 --- a/tests/kafkatest/automq/compaction_test.py +++ b/tests/kafkatest/automq/compaction_test.py @@ -1,12 +1,5 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ import time @@ -43,7 +36,7 @@ def __init__(self, test_context): self.s3_wal_upload_threshold = 50 * 1024 self.automq_stream_object_compaction_jitter_max_delay_minute = 1 - def create_kafka(self, num_nodes=1, partition=1, broker_wal='file', env=None): + def create_kafka(self, num_nodes=1, partition=1, broker_wal='s3', env=None): """ Create and configure Kafka service. @@ -91,8 +84,8 @@ def create_kafka(self, num_nodes=1, partition=1, broker_wal='file', env=None): @cluster(num_nodes=4) @matrix(stream_set_object_compaction=[True, False], - stream_object_compaction_type=[STREAM_OBJECT_COMPACTION_TYPE_MINOR_V1, STREAM_OBJECT_COMPACTION_TYPE_MAJOR_V1], wal=['file', 's3']) - @matrix(stream_set_object_compaction=[True], stream_object_compaction_type=['None'], wal=['file', 's3']) + stream_object_compaction_type=[STREAM_OBJECT_COMPACTION_TYPE_MINOR_V1, STREAM_OBJECT_COMPACTION_TYPE_MAJOR_V1], wal=['s3']) + @matrix(stream_set_object_compaction=[True], stream_object_compaction_type=['None'], wal=['s3']) def test_case(self, stream_set_object_compaction, stream_object_compaction_type, wal): ''' @@ -104,7 +97,7 @@ def test_case(self, stream_set_object_compaction, stream_object_compaction_type, self.run0(stream_set_object_compaction, stream_object_compaction_type, wal) def run0(self, stream_set_object_compaction=False, - stream_object_compaction_type=STREAM_OBJECT_COMPACTION_TYPE_MINOR_V1, wal=FILE_WAL): + stream_object_compaction_type=STREAM_OBJECT_COMPACTION_TYPE_MINOR_V1, wal=S3_WAL): """ Run the test with specified compaction type. diff --git a/tests/kafkatest/automq/memory_occupancy_test.py b/tests/kafkatest/automq/memory_occupancy_test.py index 093d2376bb..899b78db1a 100644 --- a/tests/kafkatest/automq/memory_occupancy_test.py +++ b/tests/kafkatest/automq/memory_occupancy_test.py @@ -1,12 +1,5 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ import time @@ -32,7 +25,7 @@ def __init__(self, test_context): self.consume_group = 'test_group' self.records_consumed = [] - def create_kafka(self, num_nodes=1, partition=None, log_size=None, block_size=None, wal='file', **kwargs): + def create_kafka(self, num_nodes=1, partition=None, log_size=None, block_size=None, wal='s3', **kwargs): """ Create and configure Kafka service. @@ -103,7 +96,7 @@ def check_the_consumption_quantity(self, records): assert int(receive_num) == records, f"Receive count does not match the expected records count: expected {records}, but got {receive_num}" @cluster(num_nodes=3) - @matrix(partition=[128, 512], log_size=[256 * 1024 * 1024], block_size=[128 * 1024 * 1024, 256 * 1024 * 1024], wal=['file', 's3']) + @matrix(partition=[128, 512], log_size=[256 * 1024 * 1024], block_size=[128 * 1024 * 1024, 256 * 1024 * 1024], wal=['s3']) def test(self, partition, log_size, block_size, wal): """ At any time, 1/writable record in Metric<=log cache size+100MB diff --git a/tests/kafkatest/automq/quota_test.py b/tests/kafkatest/automq/quota_test.py index bbf04b74d3..f863e22b4a 100644 --- a/tests/kafkatest/automq/quota_test.py +++ b/tests/kafkatest/automq/quota_test.py @@ -1,12 +1,5 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ from ducktape.mark.resource import cluster @@ -98,7 +91,7 @@ def start_console_consumer(self): assert len(messages) > 0, "consumer %d didn't consume any message before timeout" % idx @cluster(num_nodes=5) - @matrix(broker_in=[2500000], broker_out=[2000000], wal=['file', 's3']) + @matrix(broker_in=[2500000], broker_out=[2000000], wal=['s3']) def test_quota(self, broker_in, broker_out, wal): self.create_kafka(self.test_context, broker_in, broker_out, wal) self.kafka.start() diff --git a/tests/kafkatest/automq/s3_leakage_test.py b/tests/kafkatest/automq/s3_leakage_test.py index 43f4e46eb3..63ddbc9190 100644 --- a/tests/kafkatest/automq/s3_leakage_test.py +++ b/tests/kafkatest/automq/s3_leakage_test.py @@ -1,12 +1,5 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ import time @@ -43,7 +36,7 @@ def __init__(self, test_context): self.s3_wal_upload_threshold = 16 * 1024 * 1024 self.automq_stream_object_compaction_jitter_max_delay = 1 - def create_kafka(self, num_nodes=1, partition=1, broker_wal='file', env=None): + def create_kafka(self, num_nodes=1, partition=1, broker_wal='s3', env=None): """ Create and configure Kafka service. @@ -150,7 +143,7 @@ def run0(self, stream_object_compaction_type, wal, env=None): ) @cluster(num_nodes=2) - @matrix(wal=['file', 's3']) + @matrix(wal=['s3']) def test_s3_leak_major_v1(self, wal): """ Test S3 leak with major V1 compaction. @@ -160,7 +153,7 @@ def test_s3_leak_major_v1(self, wal): self.run0(stream_object_compaction_type=STREAM_OBJECT_COMPACTION_TYPE_MAJOR_V1, wal=wal, env=[f'AUTOMQ_STREAM_COMPACTION_MINOR_V1_COMPACTION_SIZE_THRESHOLD=0']) @cluster(num_nodes=2) - @matrix(wal=['file', 's3']) + @matrix(wal=['s3']) def test_s3_leak_minor_v1(self, wal): """ Test S3 leak with minor V1 compaction. diff --git a/tests/kafkatest/automq/version_upgrade_test.py b/tests/kafkatest/automq/version_upgrade_test.py index 59591a59d0..7016b93b52 100644 --- a/tests/kafkatest/automq/version_upgrade_test.py +++ b/tests/kafkatest/automq/version_upgrade_test.py @@ -1,12 +1,5 @@ """ -Copyright 2024, AutoMQ HK Limited. - -The use of this file is governed by the Business Source License, -as detailed in the file "/LICENSE.S3Stream" included in this repository. - -As of the Change Date specified in that file, in accordance with -the Business Source License, use of this software will be governed -by the Apache License, Version 2.0 +Copyright 2025, AutoMQ HK Limited. Licensed under Apache-2.0. """ import time diff --git a/tests/kafkatest/benchmarks/core/benchmark_test.py b/tests/kafkatest/benchmarks/core/benchmark_test.py index 321ba6e8be..e21d39addc 100644 --- a/tests/kafkatest/benchmarks/core/benchmark_test.py +++ b/tests/kafkatest/benchmarks/core/benchmark_test.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ducktape.mark import matrix +from ducktape.mark import matrix, ignore from ducktape.mark import parametrize from ducktape.mark.resource import cluster from ducktape.services.service import Service @@ -158,6 +158,7 @@ def test_long_term_producer_throughput(self, compression_type="none", self.logger.info("\n".join(summary)) return data + @ignore @cluster(num_nodes=5) @matrix(security_protocol=['SSL'], interbroker_security_protocol=['PLAINTEXT'], tls_version=['TLSv1.2', 'TLSv1.3'], compression_type=["none", "snappy"]) @matrix(security_protocol=['PLAINTEXT'], compression_type=["none", "snappy"]) diff --git a/tests/kafkatest/sanity_checks/test_performance_services.py b/tests/kafkatest/sanity_checks/test_performance_services.py index f54922991f..d94c6ae4b9 100644 --- a/tests/kafkatest/sanity_checks/test_performance_services.py +++ b/tests/kafkatest/sanity_checks/test_performance_services.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ducktape.mark import matrix, parametrize +from ducktape.mark import matrix, parametrize, ignore from ducktape.mark.resource import cluster from ducktape.tests.test import Test @@ -36,7 +36,7 @@ def __init__(self, test_context): def setUp(self): if self.zk: self.zk.start() - + @ignore #test performance send and receive in autoMQ is meaningless @cluster(num_nodes=5) # We are keeping 0.8.2 here so that we don't inadvertently break support for it. Since this is just a sanity check, # the overhead should be manageable. diff --git a/tests/kafkatest/services/connect.py b/tests/kafkatest/services/connect.py index c84a3ec43c..4ef9c4000c 100644 --- a/tests/kafkatest/services/connect.py +++ b/tests/kafkatest/services/connect.py @@ -79,6 +79,7 @@ def __init__(self, context, num_nodes, kafka, files, startup_timeout_sec=60, self.startup_timeout_sec = startup_timeout_sec self.environment = {} self.external_config_template_func = None + self.connector_config_templates = [] self.include_filestream_connectors = include_filestream_connectors self.logger.debug("include_filestream_connectors % s", include_filestream_connectors) diff --git a/tests/kafkatest/services/external_services.py b/tests/kafkatest/services/external_services.py new file mode 100644 index 0000000000..78fd8c8b6f --- /dev/null +++ b/tests/kafkatest/services/external_services.py @@ -0,0 +1,63 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import subprocess + +class DockerComposeService: + """ + A helper class to manage the lifecycle of an external service + defined in a docker-compose.yaml file. + This is NOT a ducktape service and must be managed manually from test code. + """ + def __init__(self, compose_file_path, logger): + """ + :param compose_file_path: Path to the docker-compose.yaml file. + :param logger: The test logger instance. + """ + self.compose_file_path = compose_file_path + self.logger = logger + + def start(self, env=None): + """ + Starts the service using 'docker compose up'. + :param env: A dictionary of environment variables to pass to the command. + """ + self.logger.info(f"Manually starting external service from {self.compose_file_path}...") + self._run_command("up -d", env) + + def stop(self): + """ + Stops the service using 'docker compose down'. + """ + self.logger.info(f"Manually stopping external service from {self.compose_file_path}...") + self._run_command("down --remove-orphans -v") + + def _run_command(self, command, env=None): + env_prefix = "" + if env: + for key, value in env.items(): + env_prefix += f"{key}='{value}' " + + # Use sudo -E to preserve environment variables for the docker compose command. + cmd = f"{env_prefix} sudo -E docker compose -f {self.compose_file_path} {command}" + + try: + self.logger.info(f"Running command: {cmd}") + subprocess.check_call(cmd, shell=True) + except subprocess.CalledProcessError as e: + self.logger.error(f"Failed to run command: {cmd}. Error: {e}") + log_cmd = f"{env_prefix} sudo docker compose -f {self.compose_file_path} logs" + subprocess.run(log_cmd, shell=True) + raise \ No newline at end of file diff --git a/tests/kafkatest/services/kafka/templates/kafka.properties b/tests/kafkatest/services/kafka/templates/kafka.properties index 04803bb3da..0cd1fcd061 100644 --- a/tests/kafkatest/services/kafka/templates/kafka.properties +++ b/tests/kafkatest/services/kafka/templates/kafka.properties @@ -138,10 +138,9 @@ autobalancer.controller.metrics.delay.ms=20000 elasticstream.enable=true elasticstream.endpoint=s3:// -s3.endpoint=http://10.5.0.2:4566 -s3.region=us-east-1 -s3.bucket=ko3 -s3.wal.path=/mnt/kafka/s3wal +s3.wal.path=0@s3://ko3?region=us-east-1&endpoint=http://10.5.0.2:4566&batchInterval=50 +s3.data.buckets=0@s3://ko3?region=us-east-1&endpoint=http://10.5.0.2:4566 +s3.ops.buckets=0@s3://ko3?region=us-east-1&endpoint=http://10.5.0.2:4566 s3.wal.capacity=209715200 s3.wal.cache.size=104857600 s3.wal.upload.threshold=52428800 diff --git a/tests/kafkatest/services/kafka/templates/log4j.properties b/tests/kafkatest/services/kafka/templates/log4j.properties index e37b3b7af7..0653bc7c74 100644 --- a/tests/kafkatest/services/kafka/templates/log4j.properties +++ b/tests/kafkatest/services/kafka/templates/log4j.properties @@ -20,42 +20,42 @@ log4j.appender.stdout.layout=org.apache.log4j.PatternLayout log4j.appender.stdout.layout.ConversionPattern=[%d] %p %m (%c)%n # INFO level appenders -log4j.appender.kafkaInfoAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.kafkaInfoAppender=com.automq.log.S3RollingFileAppender log4j.appender.kafkaInfoAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.kafkaInfoAppender.File={{ log_dir }}/info/server.log log4j.appender.kafkaInfoAppender.layout=org.apache.log4j.PatternLayout log4j.appender.kafkaInfoAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.kafkaInfoAppender.Threshold=INFO -log4j.appender.stateChangeInfoAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.stateChangeInfoAppender=com.automq.log.S3RollingFileAppender log4j.appender.stateChangeInfoAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.stateChangeInfoAppender.File={{ log_dir }}/info/state-change.log log4j.appender.stateChangeInfoAppender.layout=org.apache.log4j.PatternLayout log4j.appender.stateChangeInfoAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.stateChangeInfoAppender.Threshold=INFO -log4j.appender.requestInfoAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.requestInfoAppender=com.automq.log.S3RollingFileAppender log4j.appender.requestInfoAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.requestInfoAppender.File={{ log_dir }}/info/kafka-request.log log4j.appender.requestInfoAppender.layout=org.apache.log4j.PatternLayout log4j.appender.requestInfoAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.requestInfoAppender.Threshold=INFO -log4j.appender.cleanerInfoAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.cleanerInfoAppender=com.automq.log.S3RollingFileAppender log4j.appender.cleanerInfoAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.cleanerInfoAppender.File={{ log_dir }}/info/log-cleaner.log log4j.appender.cleanerInfoAppender.layout=org.apache.log4j.PatternLayout log4j.appender.cleanerInfoAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.cleanerInfoAppender.Threshold=INFO -log4j.appender.controllerInfoAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.controllerInfoAppender=com.automq.log.S3RollingFileAppender log4j.appender.controllerInfoAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.controllerInfoAppender.File={{ log_dir }}/info/controller.log log4j.appender.controllerInfoAppender.layout=org.apache.log4j.PatternLayout log4j.appender.controllerInfoAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.controllerInfoAppender.Threshold=INFO -log4j.appender.authorizerInfoAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.authorizerInfoAppender=com.automq.log.S3RollingFileAppender log4j.appender.authorizerInfoAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.authorizerInfoAppender.File={{ log_dir }}/info/kafka-authorizer.log log4j.appender.authorizerInfoAppender.layout=org.apache.log4j.PatternLayout @@ -63,49 +63,49 @@ log4j.appender.authorizerInfoAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.authorizerInfoAppender.Threshold=INFO # DEBUG level appenders -log4j.appender.kafkaDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.kafkaDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.kafkaDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.kafkaDebugAppender.File={{ log_dir }}/debug/server.log log4j.appender.kafkaDebugAppender.layout=org.apache.log4j.PatternLayout log4j.appender.kafkaDebugAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.kafkaDebugAppender.Threshold=DEBUG -log4j.appender.stateChangeDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.stateChangeDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.stateChangeDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.stateChangeDebugAppender.File={{ log_dir }}/debug/state-change.log log4j.appender.stateChangeDebugAppender.layout=org.apache.log4j.PatternLayout log4j.appender.stateChangeDebugAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.stateChangeDebugAppender.Threshold=DEBUG -log4j.appender.requestDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.requestDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.requestDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.requestDebugAppender.File={{ log_dir }}/debug/kafka-request.log log4j.appender.requestDebugAppender.layout=org.apache.log4j.PatternLayout log4j.appender.requestDebugAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.requestDebugAppender.Threshold=DEBUG -log4j.appender.cleanerDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.cleanerDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.cleanerDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.cleanerDebugAppender.File={{ log_dir }}/debug/log-cleaner.log log4j.appender.cleanerDebugAppender.layout=org.apache.log4j.PatternLayout log4j.appender.cleanerDebugAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.cleanerDebugAppender.Threshold=DEBUG -log4j.appender.controllerDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.controllerDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.controllerDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.controllerDebugAppender.File={{ log_dir }}/debug/controller.log log4j.appender.controllerDebugAppender.layout=org.apache.log4j.PatternLayout log4j.appender.controllerDebugAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.controllerDebugAppender.Threshold=DEBUG -log4j.appender.authorizerDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.authorizerDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.authorizerDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.authorizerDebugAppender.File={{ log_dir }}/debug/kafka-authorizer.log log4j.appender.authorizerDebugAppender.layout=org.apache.log4j.PatternLayout log4j.appender.authorizerDebugAppender.layout.ConversionPattern=[%d] %p %m (%c)%n log4j.appender.authorizerDebugAppender.Threshold=DEBUG -log4j.appender.autoBalancerDebugAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.autoBalancerDebugAppender=com.automq.log.S3RollingFileAppender log4j.appender.autoBalancerDebugAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.autoBalancerDebugAppender.File={{ log_dir }}/info/auto-balancer.log log4j.appender.autoBalancerDebugAppender.layout=org.apache.log4j.PatternLayout @@ -114,7 +114,7 @@ log4j.appender.autoBalancerDebugAppender.Threshold=DEBUG # TRACE level appenders -log4j.appender.s3ObjectTraceAppender=org.apache.log4j.DailyRollingFileAppender +log4j.appender.s3ObjectTraceAppender=com.automq.log.S3RollingFileAppender log4j.appender.s3ObjectTraceAppender.DatePattern='.'yyyy-MM-dd-HH log4j.appender.s3ObjectTraceAppender.File={{ log_dir }}/info/s3-object.log log4j.appender.s3ObjectTraceAppender.layout=org.apache.log4j.PatternLayout diff --git a/tests/kafkatest/services/log_compaction_tester.py b/tests/kafkatest/services/log_compaction_tester.py index cc6bf4fc29..3e42ab084f 100644 --- a/tests/kafkatest/services/log_compaction_tester.py +++ b/tests/kafkatest/services/log_compaction_tester.py @@ -63,7 +63,7 @@ def start_cmd(self, node): cmd += " export CLASSPATH;" cmd += self.path.script("kafka-run-class.sh", node) cmd += " %s" % self.java_class_name() - cmd += " --bootstrap-server %s --messages 1000000 --sleep 20 --duplicates 10 --percent-deletes 10" % (self.kafka.bootstrap_servers(self.security_protocol)) + cmd += " --bootstrap-server %s --messages 1000 --sleep 20 --duplicates 10 --percent-deletes 10" % (self.kafka.bootstrap_servers(self.security_protocol)) cmd += " 2>> %s | tee -a %s &" % (self.logs["tool_logs"]["path"], self.logs["tool_logs"]["path"]) return cmd diff --git a/tests/kafkatest/services/performance/__init__.py b/tests/kafkatest/services/performance/__init__.py index 69686f7505..8c3d4e75cf 100644 --- a/tests/kafkatest/services/performance/__init__.py +++ b/tests/kafkatest/services/performance/__init__.py @@ -16,4 +16,5 @@ from .performance import PerformanceService, throughput, latency, compute_aggregate_throughput from .end_to_end_latency import EndToEndLatencyService from .producer_performance import ProducerPerformanceService +from .automq_performance import AutoMQPerformanceService from .consumer_performance import ConsumerPerformanceService diff --git a/tests/kafkatest/services/performance/automq_performance.py b/tests/kafkatest/services/performance/automq_performance.py new file mode 100644 index 0000000000..f13b6c6821 --- /dev/null +++ b/tests/kafkatest/services/performance/automq_performance.py @@ -0,0 +1,229 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import json +import time +from ducktape.utils.util import wait_until + +from kafkatest.services.monitor.http import HttpMetricsCollector +from kafkatest.services.performance import PerformanceService +from kafkatest.services.security.security_config import SecurityConfig +from kafkatest.version import DEV_BRANCH + + +class AutoMQPerformanceService(HttpMetricsCollector, PerformanceService): + """ + Wrapper to run AutoMQ PerfCommand (bin/automq-perf-test.sh) from ducktape. + Supports Avro via Schema Registry by passing --value-schema and optional --values-file. + """ + + PERSISTENT_ROOT = "/mnt/automq_perf" + STDOUT_CAPTURE = os.path.join(PERSISTENT_ROOT, "automq_perf.stdout") + STDERR_CAPTURE = os.path.join(PERSISTENT_ROOT, "automq_perf.stderr") + LOG_DIR = os.path.join(PERSISTENT_ROOT, "logs") + LOG_FILE = os.path.join(LOG_DIR, "automq_perf.log") + LOG4J_CONFIG = os.path.join(PERSISTENT_ROOT, "tools-log4j.properties") + + def __init__(self, context, num_nodes, kafka, version=DEV_BRANCH, + producers_per_topic=1, groups_per_topic=0, consumers_per_group=1, + topics=1, partitions_per_topic=1, + send_rate=1000, record_size=1024, + topic_prefix="tt", await_topic_ready=False, + topic_configs=None, producer_configs=None, consumer_configs=None, + test_duration_minutes=1, warmup_duration_minutes=0, + value_schema=None, values_file=None): + super(AutoMQPerformanceService, self).__init__(context=context, num_nodes=num_nodes) + + self.logs = { + "automq_perf_stdout": { + "path": AutoMQPerformanceService.STDOUT_CAPTURE, + "collect_default": True}, + "automq_perf_stderr": { + "path": AutoMQPerformanceService.STDERR_CAPTURE, + "collect_default": True}, + "automq_perf_log": { + "path": AutoMQPerformanceService.LOG_FILE, + "collect_default": True} + } + + self.kafka = kafka + self.security_config = kafka.security_config.client_config() + + assert version.consumer_supports_bootstrap_server() or \ + self.security_config.security_protocol == SecurityConfig.PLAINTEXT + + self.params = { + "producers_per_topic": producers_per_topic, + "groups_per_topic": groups_per_topic, + "consumers_per_group": consumers_per_group, + "topics": topics, + "partitions_per_topic": partitions_per_topic, + "send_rate": send_rate, + "record_size": record_size, + "topic_prefix": topic_prefix, + "await_topic_ready": await_topic_ready, + "test_duration_minutes": test_duration_minutes, + "warmup_duration_minutes": warmup_duration_minutes, + } + self.topic_configs = topic_configs or {} + self.producer_configs = producer_configs or {} + self.consumer_configs = consumer_configs or {} + self.value_schema = value_schema + self.values_file = values_file + + for node in self.nodes: + node.version = version + + def start_cmd(self, node): + bs = self.kafka.bootstrap_servers(self.security_config.security_protocol) + script = self.path.script("automq-perf-test.sh", node) + + args = [] + args.append(f"-B {bs}") + args.append(f"-p {self.params['producers_per_topic']}") + args.append(f"-g {self.params['groups_per_topic']}") + args.append(f"-c {self.params['consumers_per_group']}") + args.append(f"-t {self.params['topics']}") + args.append(f"-n {self.params['partitions_per_topic']}") + args.append(f"-r {self.params['send_rate']}") + args.append(f"-s {self.params['record_size']}") + args.append(f"-X {self.params['topic_prefix']}") + args.append(f"--await-topic-ready {str(self.params['await_topic_ready']).lower()}") + args.append(f"-w {self.params['warmup_duration_minutes']}") + args.append(f"-d {self.params['test_duration_minutes']}") + + # Topic/producer/consumer configs + if self.topic_configs: + topic_cfgs = " ".join([f"{k}={v}" for k, v in self.topic_configs.items()]) + args.append(f"-T {topic_cfgs}") + if self.producer_configs: + prod_cfgs = " ".join([f"{k}={v}" for k, v in self.producer_configs.items()]) + args.append(f"-P {prod_cfgs}") + if self.consumer_configs: + cons_cfgs = " ".join([f"{k}={v}" for k, v in self.consumer_configs.items()]) + args.append(f"-C {cons_cfgs}") + + # Avro schema options + if self.value_schema: + args.append(f"--value-schema '{self.value_schema}'") + if self.values_file: + args.append(f"--values-file {self.values_file}") + + # Run under PERSISTENT_ROOT so perf-*.json is created there + cmd = ( + f"cd {AutoMQPerformanceService.PERSISTENT_ROOT}; " + f"{script} {' '.join(args)} 2>>{AutoMQPerformanceService.STDERR_CAPTURE} | tee {AutoMQPerformanceService.STDOUT_CAPTURE}" + ) + return cmd + + def _worker(self, idx, node): + node.account.ssh(f"mkdir -p {AutoMQPerformanceService.PERSISTENT_ROOT}", allow_fail=False) + + cmd = self.start_cmd(node) + self.logger.debug("AutoMQ perf command: %s", cmd) + + start = time.time() + proc = node.account.ssh_capture(cmd) + first_line = next(proc, None) + if first_line is None: + raise Exception("No output from AutoMQ performance command") + + # consume output until process exits + for _ in proc: + pass + elapsed = time.time() - start + self.logger.debug("AutoMQ PerfCommand ran for %s seconds" % elapsed) + + # Try to parse result file path and produced count + result_file = None + produced = None + try: + # Prefer explicit line if present in stdout + for line in node.account.ssh_capture(f"cat {AutoMQPerformanceService.STDOUT_CAPTURE}"): + if "Saving results to" in line: + result_file = line.strip().split("Saving results to")[-1].strip() + break + # Fallback: find latest perf-*.json in CWD + if not result_file: + candidates = list(node.account.ssh_capture( + f"ls -1t {AutoMQPerformanceService.PERSISTENT_ROOT}/perf-*.json 2>/dev/null | head -1" + )) + if candidates: + result_file = candidates[0].strip() + if result_file: + # Ensure absolute path when parsed from stdout + if not result_file.startswith('/'): + result_file = f"{AutoMQPerformanceService.PERSISTENT_ROOT}/{result_file}" + # Copy for unified collection + node.account.ssh( + f"cp -f {result_file} {AutoMQPerformanceService.PERSISTENT_ROOT}/result.json || true", + allow_fail=True + ) + content = "".join(node.account.ssh_capture(f"cat {result_file}")) + # Log full content for validation and future analysis + self.logger.info(f"[AutoMQPerf] Result file: {result_file}\nContent: {content}") + data = json.loads(content) + produced = int(data.get("produceCountTotal", 0)) + except Exception: + pass + # record to results for test usage (simple, extensible structure) + if self.results is not None and len(self.results) > 0: + parsed = { + "produced": produced, + "consumed": None, + "produce_rate": None, + "consume_rate": None, + "produce_throughput_bps": None, + "consume_throughput_bps": None, + "errors": None, + "backlog_last": None, + "result_file": result_file, + } + try: + data_map = json.loads(content) if 'content' in locals() else None + if isinstance(data_map, dict): + parsed["consumed"] = int(data_map.get("consumeCountTotal", 0)) + parsed["produce_rate"] = float(data_map.get("produceRateTotal", 0.0)) + parsed["consume_rate"] = float(data_map.get("consumeRateTotal", 0.0)) + parsed["produce_throughput_bps"] = float(data_map.get("produceThroughputTotalBps", 0.0)) + parsed["consume_throughput_bps"] = float(data_map.get("consumeThroughputTotalBps", 0.0)) + parsed["errors"] = int(data_map.get("produceErrorTotal", 0)) + bl = data_map.get("backlog") + if isinstance(bl, list) and bl: + try: + parsed["backlog_last"] = int(bl[-1]) + except Exception: + parsed["backlog_last"] = None + except Exception: + pass + + self.results[idx-1] = parsed + + def stop(self, **kwargs): + """Stop the service but swallow non-fatal exceptions to keep log collection intact.""" + try: + super(AutoMQPerformanceService, self).stop(**kwargs) + except Exception as e: + try: + # Attempt to continue even if reverse forwarder or httpd shutdown races + self.logger.warn(f"AutoMQPerf stop encountered error: {e}") + except Exception: + pass + + def java_class_name(self): + # Target the actual perf runner class to avoid killing unrelated Java processes + # See bin/automq-perf-test.sh + return "org.apache.kafka.tools.automq.PerfCommand" diff --git a/tests/kafkatest/services/templates/connect_log4j.properties b/tests/kafkatest/services/templates/connect_log4j.properties index 90543ac946..e8730cfb8e 100644 --- a/tests/kafkatest/services/templates/connect_log4j.properties +++ b/tests/kafkatest/services/templates/connect_log4j.properties @@ -16,13 +16,31 @@ ## # Define the root logger with appender file -log4j.rootLogger = {{ log_level|default("INFO") }}, FILE +log4j.rootLogger=INFO, stdout, connectAppender +# Send the logs to the console. +# +log4j.appender.stdout=org.apache.log4j.ConsoleAppender +log4j.appender.stdout.layout=org.apache.log4j.PatternLayout + +# Send the logs to a file, rolling the file when it reaches the specified size. For example, the `File` option specifies the +# location of the log files (e.g. ${kafka.logs.dir}/connect.log). The `MaxFileSize` option specifies the maximum size of the log file, +# and the `MaxBackupIndex` option specifies the number of backup files to keep. +# +log4j.appender.connectAppender=com.automq.log.S3RollingFileAppender +log4j.appender.connectAppender.configProviderClass=org.apache.kafka.connect.automq.log.ConnectS3LogConfigProvider +log4j.appender.connectAppender.MaxFileSize=10MB +log4j.appender.connectAppender.MaxBackupIndex=11 +log4j.appender.connectAppender.File={{ log_file }} +log4j.appender.connectAppender.layout=org.apache.log4j.PatternLayout -log4j.appender.FILE=org.apache.log4j.FileAppender -log4j.appender.FILE.File={{ log_file }} -log4j.appender.FILE.ImmediateFlush=true -log4j.appender.FILE.Append=true -log4j.appender.FILE.layout=org.apache.log4j.PatternLayout -log4j.appender.FILE.layout.conversionPattern=[%d] %p %m (%c)%n +# The `%X{connector.context}` parameter in the layout includes connector-specific and task-specific information +# in the log messages, where appropriate. This makes it easier to identify those log messages that apply to a +# specific connector. +# +connect.log.pattern=[%d] %p %X{connector.context}%m (%c:%L)%n + +log4j.appender.stdout.layout.ConversionPattern=${connect.log.pattern} +log4j.appender.connectAppender.layout.ConversionPattern=${connect.log.pattern} log4j.logger.org.reflections=ERROR + diff --git a/tests/kafkatest/tests/client/consumer_test.py b/tests/kafkatest/tests/client/consumer_test.py index 7cb66069a7..163d848af6 100644 --- a/tests/kafkatest/tests/client/consumer_test.py +++ b/tests/kafkatest/tests/client/consumer_test.py @@ -83,8 +83,6 @@ def setup_consumer(self, topic, **kwargs): metadata_quorum=[quorum.isolated_kraft], use_new_coordinator=[True], group_protocol=consumer_group.all_group_protocols - use_new_coordinator=[True], - group_protocol=consumer_group.all_group_protocols ) def test_broker_rolling_bounce(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None): def test_broker_rolling_bounce(self, metadata_quorum=quorum.zk, use_new_coordinator=False, group_protocol=None): diff --git a/tests/kafkatest/tests/connect/connect_az_metadata_test.py b/tests/kafkatest/tests/connect/connect_az_metadata_test.py new file mode 100644 index 0000000000..c67e83ecdd --- /dev/null +++ b/tests/kafkatest/tests/connect/connect_az_metadata_test.py @@ -0,0 +1,210 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import textwrap + +from ducktape.mark.resource import cluster +from ducktape.utils.util import wait_until + +from kafkatest.tests.kafka_test import KafkaTest +from kafkatest.services.connect import ConnectDistributedService, VerifiableSink, VerifiableSource + + +class ConnectAzMetadataTest(KafkaTest): + """End-to-end validation that Connect honors AzMetadataProvider when building client configs.""" + + AZ_CONFIG_KEY = "automq.test.az.id" + EXPECTED_AZ = "test-az-1" + TOPIC = "az-aware-connect" + FILE_SOURCE_CONNECTOR = 'org.apache.kafka.connect.file.FileStreamSourceConnector' + FILE_SINK_CONNECTOR = 'org.apache.kafka.connect.file.FileStreamSinkConnector' + + INPUT_FILE = "/mnt/connect.input" + OUTPUT_FILE = "/mnt/connect.output" + + OFFSETS_TOPIC = "connect-offsets" + OFFSETS_REPLICATION_FACTOR = "1" + OFFSETS_PARTITIONS = "1" + CONFIG_TOPIC = "connect-configs" + CONFIG_REPLICATION_FACTOR = "1" + STATUS_TOPIC = "connect-status" + STATUS_REPLICATION_FACTOR = "1" + STATUS_PARTITIONS = "1" + EXACTLY_ONCE_SOURCE_SUPPORT = "disabled" + SCHEDULED_REBALANCE_MAX_DELAY_MS = "60000" + CONNECT_PROTOCOL="sessioned" + + # Since tasks can be assigned to any node and we're testing with files, we need to make sure the content is the same + # across all nodes. + FIRST_INPUT_LIST = ["foo", "bar", "baz"] + FIRST_INPUTS = "\n".join(FIRST_INPUT_LIST) + "\n" + SECOND_INPUT_LIST = ["razz", "ma", "tazz"] + SECOND_INPUTS = "\n".join(SECOND_INPUT_LIST) + "\n" + + SCHEMA = { "type": "string", "optional": False } + + def __init__(self, test_context): + super(ConnectAzMetadataTest, self).__init__(test_context, num_zk=1, num_brokers=1) + # Single worker is sufficient for testing AZ metadata provider + self.cc = ConnectDistributedService(test_context, 1, self.kafka, []) + self.source = None + self.sink = None + self._last_describe_output = "" + + @cluster(num_nodes=4) + def test_consumer_metadata_contains_az(self): + if self.zk: + self.zk.start() + self.kafka.start() + + self.cc.clean() + self._install_az_provider_plugin() + + self.cc.set_configs(lambda node: self._render_worker_config(node)) + self.cc.start() + + try: + self.source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=50) + self.source.start() + wait_until(lambda: len(self.source.sent_messages()) > 0, timeout_sec=60, + err_msg="Timed out waiting for VerifiableSource to emit records") + + self.sink = VerifiableSink(self.cc, topics=[self.TOPIC]) + self.sink.start() + wait_until(lambda: len(self.sink.received_messages()) > 0, timeout_sec=60, + err_msg="Timed out waiting for VerifiableSink to consume records") + + group_id = "connect-%s" % self.sink.name + + def az_metadata_present(): + output = self.kafka.describe_consumer_group(group_id) + self._last_describe_output = output + return self._consumer_group_has_expected_metadata(output) + + wait_until(az_metadata_present, timeout_sec=60, + err_msg="Consumer group metadata never reflected AZ-aware client settings") + + # Final verification that AZ metadata is present + assert self._consumer_group_has_expected_metadata(self._last_describe_output), \ + "Final consumer group output did not contain expected AZ metadata: %s" % self._last_describe_output + finally: + if self.sink is not None: + self.sink.stop() + if self.source is not None: + self.source.stop() + self.cc.stop() + + def _render_worker_config(self, node): + base_config = self.render("connect-distributed.properties", node=node) + # Ensure the worker passes the AZ hint down to the ServiceLoader plugin + return base_config + "\n%s=%s\n" % (self.AZ_CONFIG_KEY, self.EXPECTED_AZ) + + def _install_az_provider_plugin(self): + # Create a simple mock AzMetadataProvider implementation directly in the Connect runtime classpath + java_source = textwrap.dedent(""" + package org.apache.kafka.connect.automq.test; + + import java.util.Map; + import java.util.Optional; + import org.apache.kafka.connect.automq.az.AzMetadataProvider; + + public class FixedAzMetadataProvider implements AzMetadataProvider {{ + private volatile Optional availabilityZoneId = Optional.empty(); + + @Override + public void configure(Map workerProps) {{ + System.out.println("FixedAzMetadataProvider.configure() called with worker properties: " + workerProps.keySet()); + + String az = workerProps.get("{}"); + System.out.println("AZ config value for key '{}': " + az); + + if (az == null || az.isBlank()) {{ + availabilityZoneId = Optional.empty(); + System.out.println("FixedAzMetadataProvider: No AZ configured, setting to empty"); + }} else {{ + availabilityZoneId = Optional.of(az); + System.out.println("FixedAzMetadataProvider: Setting AZ to: " + az); + }} + }} + + @Override + public Optional availabilityZoneId() {{ + System.out.println("FixedAzMetadataProvider.availabilityZoneId() called, returning: " + availabilityZoneId.orElse("empty")); + return availabilityZoneId; + }} + }} + """.format(self.AZ_CONFIG_KEY, self.AZ_CONFIG_KEY)) + + service_definition = "org.apache.kafka.connect.automq.test.FixedAzMetadataProvider\n" + + for node in self.cc.nodes: + # Get the Connect runtime classes directory where ServiceLoader will find our class + kafka_home = self.cc.path.home() + runtime_classes_dir = f"{kafka_home}/connect/runtime/build/classes/java/main" + + # Create the package directory structure in the runtime classes + test_package_dir = f"{runtime_classes_dir}/org/apache/kafka/connect/automq/test" + node.account.ssh(f"mkdir -p {test_package_dir}") + node.account.ssh(f"mkdir -p {runtime_classes_dir}/META-INF/services") + + # Write the Java source file to a temporary location + temp_src_dir = f"/tmp/az-provider-src/org/apache/kafka/connect/automq/test" + node.account.ssh(f"mkdir -p {temp_src_dir}") + java_path = f"{temp_src_dir}/FixedAzMetadataProvider.java" + node.account.create_file(java_path, java_source) + + # Create the ServiceLoader service definition in the runtime classes + service_path = f"{runtime_classes_dir}/META-INF/services/org.apache.kafka.connect.automq.az.AzMetadataProvider" + node.account.create_file(service_path, service_definition) + + # Compile the Java file directly to the runtime classes directory + classpath = f"{kafka_home}/connect/runtime/build/libs/*:{kafka_home}/connect/runtime/build/dependant-libs/*:{kafka_home}/clients/build/libs/*" + compile_cmd = f"javac -cp \"{classpath}\" -d {runtime_classes_dir} {java_path}" + print(f"Compiling with command: {compile_cmd}") + result = node.account.ssh(compile_cmd, allow_fail=False) + print(f"Compilation result: {result}") + + # Verify the compiled class exists in the runtime classes directory + class_path = f"{test_package_dir}/FixedAzMetadataProvider.class" + verify_cmd = f"ls -la {class_path}" + verify_result = node.account.ssh(verify_cmd, allow_fail=True) + print(f"Class file verification: {verify_result}") + + # Also verify the service definition exists + service_verify_cmd = f"cat {service_path}" + service_verify_result = node.account.ssh(service_verify_cmd, allow_fail=True) + print(f"Service definition verification: {service_verify_result}") + + print(f"AZ metadata provider plugin installed in runtime classpath for node {node.account.hostname}") + + def _consumer_group_has_expected_metadata(self, describe_output): + # Simply check if any line in the output contains our expected AZ metadata + # This is more robust than trying to parse the exact table format + expected_az_in_client_id = "automq_az={}".format(self.EXPECTED_AZ) + + # Debug: print the output to see what we're actually getting + print("=== Consumer Group Describe Output ===") + print(describe_output) + print("=== Looking for: {} ===".format(expected_az_in_client_id)) + + # Check if any line contains the expected AZ metadata + for line in describe_output.splitlines(): + if expected_az_in_client_id in line: + print("Found AZ metadata in line: {}".format(line)) + return True + + print("AZ metadata not found in consumer group output") + return False diff --git a/tests/kafkatest/tests/connect/connect_distributed_test.py b/tests/kafkatest/tests/connect/connect_distributed_test.py index c7f5e3a5a8..5885e60f5e 100644 --- a/tests/kafkatest/tests/connect/connect_distributed_test.py +++ b/tests/kafkatest/tests/connect/connect_distributed_test.py @@ -368,12 +368,7 @@ def test_pause_and_resume_source(self, exactly_once_source, connect_protocol, me @matrix( connect_protocol=['sessioned', 'compatible', 'eager'], metadata_quorum=[quorum.isolated_kraft], - use_new_coordinator=[False] - ) - @matrix( - connect_protocol=['sessioned', 'compatible', 'eager'], - metadata_quorum=[quorum.isolated_kraft], - use_new_coordinator=[True], + use_new_coordinator=[True, False], group_protocol=consumer_group.all_group_protocols ) def test_pause_and_resume_sink(self, connect_protocol, metadata_quorum, use_new_coordinator=False, group_protocol=None): @@ -419,7 +414,7 @@ def test_pause_and_resume_sink(self, connect_protocol, metadata_quorum, use_new_ err_msg="Failed to see connector transition to the RUNNING state") # after resuming, we should see records consumed again - wait_until(lambda: len(self.sink.received_messages()) > num_messages, timeout_sec=40, + wait_until(lambda: len(self.sink.received_messages()) > num_messages, timeout_sec=300, err_msg="Failed to consume messages after resuming sink connector") @cluster(num_nodes=5) @@ -649,7 +644,7 @@ def _wait_for_loggers(self, level, request_time, namespace, workers=None): ) @matrix( security_protocol=[SecurityConfig.PLAINTEXT, SecurityConfig.SASL_SSL], - exactly_once_source=[True, False], + exactly_once_source=[True, False], connect_protocol=['sessioned', 'compatible', 'eager'], metadata_quorum=[quorum.isolated_kraft], use_new_coordinator=[True], @@ -680,7 +675,7 @@ def test_file_source_and_sink(self, security_protocol, exactly_once_source, conn # do rebalancing of the group, etc, and b) without explicit leave group support, rebalancing takes awhile for node in self.cc.nodes: node.account.ssh("echo -e -n " + repr(self.FIRST_INPUTS) + " >> " + self.INPUT_FILE) - wait_until(lambda: self._validate_file_output(self.FIRST_INPUT_LIST), timeout_sec=90, err_msg="Data added to input file was not seen in the output file in a reasonable amount of time.") + wait_until(lambda: self._validate_file_output(self.FIRST_INPUT_LIST), timeout_sec=300, err_msg="Data added to input file was not seen in the output file in a reasonable amount of time.") # Restarting both should result in them picking up where they left off, # only processing new data. @@ -688,7 +683,7 @@ def test_file_source_and_sink(self, security_protocol, exactly_once_source, conn for node in self.cc.nodes: node.account.ssh("echo -e -n " + repr(self.SECOND_INPUTS) + " >> " + self.INPUT_FILE) - wait_until(lambda: self._validate_file_output(self.FIRST_INPUT_LIST + self.SECOND_INPUT_LIST), timeout_sec=150, err_msg="Sink output file never converged to the same state as the input file") + wait_until(lambda: self._validate_file_output(self.FIRST_INPUT_LIST + self.SECOND_INPUT_LIST), timeout_sec=300, err_msg="Sink output file never converged to the same state as the input file") @cluster(num_nodes=6) @matrix( @@ -1092,4 +1087,4 @@ def _restart_worker(self, node, clean=True): self.cc.start_node(node) monitor.wait_until("Starting connectors and tasks using config offset", timeout_sec=90, err_msg="Kafka Connect worker didn't successfully join group and start work") - self.logger.info("Bounced Kafka Connect on %s and rejoined in %f seconds", node.account, time.time() - started) + self.logger.info("Bounced Kafka Connect on %s and rejoined in %f seconds", node.account, time.time() - started) \ No newline at end of file diff --git a/tests/kafkatest/tests/connect/connect_metrics_log_test.py b/tests/kafkatest/tests/connect/connect_metrics_log_test.py new file mode 100644 index 0000000000..92525599b3 --- /dev/null +++ b/tests/kafkatest/tests/connect/connect_metrics_log_test.py @@ -0,0 +1,764 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ducktape.tests.test import Test +from ducktape.mark.resource import cluster +from ducktape.utils.util import wait_until +from ducktape.mark import matrix, parametrize + +from kafkatest.services.zookeeper import ZookeeperService +from kafkatest.services.kafka import KafkaService, config_property, quorum +from kafkatest.services.connect import ConnectDistributedService, VerifiableSource +from kafkatest.services.security.security_config import SecurityConfig +from kafkatest.version import DEV_BRANCH + +import time +import subprocess + + +class ConnectMetricsLogTest(Test): + """ + Test class specifically for testing Kafka Connect metrics endpoint functionality, + extracted from ConnectDistributedTest to focus on metrics validation. + """ + + FILE_SOURCE_CONNECTOR = 'org.apache.kafka.connect.file.FileStreamSourceConnector' + FILE_SINK_CONNECTOR = 'org.apache.kafka.connect.file.FileStreamSinkConnector' + + INPUT_FILE = "/mnt/connect.input" + OUTPUT_FILE = "/mnt/connect.output" + + TOPIC = "test" + OFFSETS_TOPIC = "connect-offsets" + OFFSETS_REPLICATION_FACTOR = "1" + OFFSETS_PARTITIONS = "1" + CONFIG_TOPIC = "connect-configs" + CONFIG_REPLICATION_FACTOR = "1" + STATUS_TOPIC = "connect-status" + STATUS_REPLICATION_FACTOR = "1" + STATUS_PARTITIONS = "1" + EXACTLY_ONCE_SOURCE_SUPPORT = "disabled" + SCHEDULED_REBALANCE_MAX_DELAY_MS = "60000" + CONNECT_PROTOCOL="sessioned" + + # Since tasks can be assigned to any node and we're testing with files, we need to make sure the content is the same + # across all nodes. + FIRST_INPUT_LIST = ["foo", "bar", "baz"] + FIRST_INPUTS = "\n".join(FIRST_INPUT_LIST) + "\n" + SECOND_INPUT_LIST = ["razz", "ma", "tazz"] + SECOND_INPUTS = "\n".join(SECOND_INPUT_LIST) + "\n" + + SCHEMA = { "type": "string", "optional": False } + + def __init__(self, test_context): + super(ConnectMetricsLogTest, self).__init__(test_context) + self.num_zk = 1 + self.num_brokers = 1 + self.topics = { + 'metrics-test-topic': {'partitions': 1, 'replication-factor': 1}, + 'test': {'partitions': 1, 'replication-factor': 1} + } + + # Constants from original test class + self.TOPIC = "test" + + self.zk = ZookeeperService(test_context, self.num_zk) if quorum.for_test(test_context) == quorum.zk else None + + def setup_services(self, + security_protocol=SecurityConfig.PLAINTEXT, + broker_version=DEV_BRANCH, + auto_create_topics=True, + num_workers=3, + kraft=True): + """Setup Kafka and Connect services""" + self.kafka = KafkaService(self.test_context, self.num_brokers, self.zk, + security_protocol=security_protocol, + interbroker_security_protocol=security_protocol, + topics=self.topics, version=broker_version, + server_prop_overrides=[ + ["auto.create.topics.enable", str(auto_create_topics)], + ["transaction.state.log.replication.factor", str(self.num_brokers)], + ["transaction.state.log.min.isr", str(self.num_brokers)] + ], allow_zk_with_kraft=kraft) + + self.cc = ConnectDistributedService(self.test_context, num_workers, self.kafka, []) + self.cc.log_level = "DEBUG" + + if self.zk: + self.zk.start() + self.kafka.start() + + def is_running(self, connector, node=None): + """Check if a connector and all its tasks are running""" + status = self._connector_status(connector.name, node) + return self._connector_has_state(status, 'RUNNING') and self._all_tasks_have_state(status, connector.tasks, 'RUNNING') + + def _connector_status(self, connector, node=None): + """Get connector status""" + try: + return self.cc.get_connector_status(connector, node) + except: + return None + + def _connector_has_state(self, status, state): + """Check if connector has specific state""" + return status is not None and status['connector']['state'] == state + + def _all_tasks_have_state(self, status, task_count, state): + """Check if all tasks have specific state""" + if status is None: + return False + + tasks = status['tasks'] + if len(tasks) != task_count: + return False + + return all(task['state'] == state for task in tasks) + + def _wait_for_metrics_available(self, timeout_sec=60): + """Wait for metrics endpoint to become available""" + self.logger.info("Waiting for metrics endpoint to become available...") + + def metrics_available(): + for node in self.cc.nodes: + try: + cmd = "curl -s http://localhost:9464/metrics" + result = node.account.ssh_capture(cmd, allow_fail=True) + metrics_output = "".join([line for line in result]) + + # Check for any metrics output (not just kafka_connect) + if len(metrics_output.strip()) > 0 and ("#" in metrics_output or "_" in metrics_output): + self.logger.info(f"Metrics available on node {node.account.hostname}, content length: {len(metrics_output)}") + return True + else: + self.logger.debug(f"Node {node.account.hostname} metrics not ready yet, output length: {len(metrics_output)}") + except Exception as e: + self.logger.debug(f"Error checking metrics on node {node.account.hostname}: {e}") + continue + return False + + wait_until( + metrics_available, + timeout_sec=timeout_sec, + err_msg="Metrics endpoint did not become available within the specified time" + ) + + self.logger.info("Metrics endpoint is now available!") + + def _verify_opentelemetry_metrics(self): + """Verify OpenTelemetry metrics content""" + for node in self.cc.nodes: + cmd = "curl -s http://localhost:9464/metrics" + result = node.account.ssh_capture(cmd) + metrics_output = "".join([line for line in result]) + + # Basic check - verify any metrics output exists + assert len(metrics_output.strip()) > 0, "Metrics endpoint returned no content" + + # Print ALL metrics for debugging + self.logger.info(f"=== ALL METRICS from Node {node.account.hostname} ===") + self.logger.info(metrics_output) + self.logger.info(f"=== END OF METRICS from Node {node.account.hostname} ===") + + # Find all metric lines (not comments) + metric_lines = [line for line in metrics_output.split('\n') + if line.strip() and not line.startswith('#') and ('_' in line or '{' in line)] + + # Should have at least some metrics + assert len(metric_lines) > 0, "No valid metric lines found" + + self.logger.info(f"Found {len(metric_lines)} metric lines") + + # Log kafka_connect metrics specifically + kafka_connect_lines = [line for line in metric_lines if 'kafka_connect' in line] + self.logger.info(f"Found {len(kafka_connect_lines)} kafka_connect metric lines:") + for i, line in enumerate(kafka_connect_lines): + self.logger.info(f"kafka_connect metric {i+1}: {line}") + + # Check for Prometheus format characteristics + has_help = "# HELP" in metrics_output + has_type = "# TYPE" in metrics_output + + if has_help and has_type: + self.logger.info("Metrics conform to Prometheus format") + else: + self.logger.warning("Metrics may not be in standard Prometheus format") + + # Use lenient metric validation to analyze values + self._validate_metric_values(metrics_output) + + self.logger.info(f"Node {node.account.hostname} basic metrics validation passed") + + def _verify_comprehensive_metrics(self): + """Comprehensive metrics validation""" + for node in self.cc.nodes: + cmd = "curl -s http://localhost:9464/metrics" + result = node.account.ssh_capture(cmd) + metrics_output = "".join([line for line in result]) + + # Basic check - verify any metrics output exists + assert len(metrics_output.strip()) > 0, "Metrics endpoint returned no content" + + # Print ALL metrics for comprehensive debugging + self.logger.info(f"=== COMPREHENSIVE METRICS from Node {node.account.hostname} ===") + self.logger.info(metrics_output) + self.logger.info(f"=== END OF COMPREHENSIVE METRICS from Node {node.account.hostname} ===") + + # Find all metric lines (start with letter, not comments) + metric_lines = [line for line in metrics_output.split('\n') + if line.strip() and not line.startswith('#') and ('_' in line or '{' in line)] + self.logger.info(f"Found metric line count: {len(metric_lines)}") + + # Find kafka_connect related metrics + kafka_connect_lines = [line for line in metric_lines if 'kafka_connect' in line] + self.logger.info(f"Found kafka_connect metric line count: {len(kafka_connect_lines)}") + + # Print all kafka_connect metrics + self.logger.info("=== ALL kafka_connect metrics ===") + for i, line in enumerate(kafka_connect_lines): + self.logger.info(f"kafka_connect metric {i+1}: {line}") + + # If no kafka_connect metrics found, show other metrics + if len(kafka_connect_lines) == 0: + self.logger.warning("No kafka_connect metrics found, showing other metrics:") + for i, line in enumerate(metric_lines[:10]): # Show first 10 instead of 5 + self.logger.info(f"Other metric line {i+1}: {line}") + + # Should have at least some metric output + assert len(metric_lines) > 0, "No valid metric lines found" + else: + # Found kafka_connect metrics + self.logger.info(f"Successfully found {len(kafka_connect_lines)} kafka_connect metrics") + + # Check for HELP and TYPE comments (Prometheus format characteristics) + has_help = "# HELP" in metrics_output + has_type = "# TYPE" in metrics_output + + if has_help: + self.logger.info("Found HELP comments - conforms to Prometheus format") + if has_type: + self.logger.info("Found TYPE comments - conforms to Prometheus format") + + self.logger.info(f"Node {node.account.hostname} metrics validation passed, total {len(metric_lines)} metrics found") + + def _validate_metric_values(self, metrics_output): + """Validate metric value reasonableness - more lenient version""" + lines = metrics_output.split('\n') + negative_metrics = [] + + self.logger.info("=== ANALYZING METRIC VALUES ===") + + for line in lines: + if line.startswith('kafka_connect_') and not line.startswith('#'): + # Parse metric line: metric_name{labels} value timestamp + parts = line.split() + if len(parts) >= 2: + try: + value = float(parts[1]) + metric_name = parts[0].split('{')[0] if '{' in parts[0] else parts[0] + + # Log all metric values for analysis + self.logger.info(f"Metric: {metric_name} = {value}") + + # Some metrics can legitimately be negative (e.g., ratios, differences, etc.) + # Only flag as problematic if it's a count or gauge that shouldn't be negative + if value < 0: + negative_metrics.append(f"{parts[0]} = {value}") + + # Allow certain metrics to be negative + allowed_negative_patterns = [ + 'ratio', + 'seconds_ago', + 'difference', + 'offset', + 'lag' + ] + + is_allowed_negative = any(pattern in parts[0].lower() for pattern in allowed_negative_patterns) + + if is_allowed_negative: + self.logger.info(f"Negative value allowed for metric: {parts[0]} = {value}") + else: + self.logger.warning(f"Potentially problematic negative value: {parts[0]} = {value}") + # Don't assert here, just log for now + + except ValueError: + # Skip unparseable lines + continue + + if negative_metrics: + self.logger.info(f"Found {len(negative_metrics)} metrics with negative values:") + for metric in negative_metrics: + self.logger.info(f" - {metric}") + + self.logger.info("=== END METRIC VALUE ANALYSIS ===") + + def _verify_metrics_updates(self): + """Verify metrics update over time""" + # Get initial metrics + initial_metrics = {} + for node in self.cc.nodes: + cmd = "curl -s http://localhost:9464/metrics" + result = node.account.ssh_capture(cmd) + initial_metrics[node] = "".join([line for line in result]) + + # Wait for some time + time.sleep(5) + + # Get metrics again and compare + for node in self.cc.nodes: + cmd = "curl -s http://localhost:9464/metrics" + result = node.account.ssh_capture(cmd) + current_metrics = "".join([line for line in result]) + + # Metrics should have changed (at least timestamps will update) + # More detailed verification can be done here + self.logger.info(f"Node {node.account.hostname} metrics have been updated") + + def _safe_cleanup(self): + """Safe resource cleanup""" + try: + # Delete connectors + connectors = self.cc.list_connectors() + for connector in connectors: + try: + self.cc.delete_connector(connector) + self.logger.info(f"Deleted connector: {connector}") + except Exception as e: + self.logger.warning(f"Failed to delete connector {connector}: {e}") + + # Stop services + self.cc.stop() + + except Exception as e: + self.logger.error(f"Error occurred during cleanup: {e}") + + def _check_port_listening(self, node, port): + """Check if a port is listening on the given node""" + try: + result = list(node.account.ssh_capture(f"netstat -ln | grep :{port}", allow_fail=True)) + return len(result) > 0 + except: + return False + + def _verify_remote_write_requests(self, node, log_file="/tmp/mock_remote_write.log"): + """Verify that remote write requests were received""" + try: + # Check the mock server log for received requests + result = list(node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + + self.logger.info(f"Remote write log content: {log_content}") + + # Look for evidence of received data + if "Received" in log_content or "received" in log_content: + self.logger.info("Remote write requests were successfully received") + return True + + # Also check if the process is running and listening + if self._check_port_listening(node, 9090) or self._check_port_listening(node, 9091): + self.logger.info("Remote write server is listening, requests may have been processed") + return True + + self.logger.warning("No clear evidence of remote write requests in log") + return False + + except Exception as e: + self.logger.warning(f"Error verifying remote write requests: {e}") + # Don't fail the test if we can't verify the log, as the server might be working + return True + + def _verify_s3_metrics_export_localstack(self, bucket_name, node, selector_type): + """Verify that metrics were exported to S3 via localstack""" + try: + # Recursively list all object files (not directories) in S3 bucket + list_cmd = f"aws s3 ls s3://{bucket_name}/ --recursive --endpoint=http://10.5.0.2:4566" + + ret, val = subprocess.getstatusoutput(list_cmd) + self.logger.info( + f'\n--------------recursive objects[bucket:{bucket_name}]--------------------\n{val}\n--------------recursive objects end--------------------\n') + if ret != 0: + self.logger.warning(f"Failed to list bucket objects recursively, return code: {ret}, output: {val}") + # Try non-recursive listing of directory structure + list_dir_cmd = f"aws s3 ls s3://{bucket_name}/ --endpoint=http://10.5.0.2:4566" + ret2, val2 = subprocess.getstatusoutput(list_dir_cmd) + self.logger.info(f"Directory listing: {val2}") + + # If non-recursive also fails, the bucket may not exist or lack permissions + if ret2 != 0: + raise Exception(f"Failed to list bucket contents, output: {val}") + else: + # Found directories but no files, upload may not be complete yet + self.logger.info("Found directories but no files yet, checking subdirectories...") + + # Try to list contents under automq/metrics/ + automq_cmd = f"aws s3 ls s3://{bucket_name}/automq/metrics/ --recursive --endpoint=http://10.5.0.2:4566" + ret3, val3 = subprocess.getstatusoutput(automq_cmd) + self.logger.info(f"AutoMQ metrics directory contents: {val3}") + + if ret3 == 0 and val3.strip(): + s3_objects = [line.strip() for line in val3.strip().split('\n') if line.strip()] + else: + return False + else: + s3_objects = [line.strip() for line in val.strip().split('\n') if line.strip()] + + self.logger.info(f"S3 bucket {bucket_name} file contents (total {len(s3_objects)} files): {s3_objects}") + + if s3_objects: + # Filter out directory lines, keep only file lines (file lines usually have size info) + file_objects = [] + for obj_line in s3_objects: + parts = obj_line.split() + # File line format: 2025-01-01 12:00:00 size_in_bytes filename + # Directory line format: PRE directory_name/ or just directory name + if len(parts) >= 4 and not obj_line.strip().startswith('PRE') and 'automq/metrics/' in obj_line: + file_objects.append(obj_line) + + self.logger.info(f"Found {len(file_objects)} actual metric files in S3:") + for file_obj in file_objects: + self.logger.info(f" - {file_obj}") + + if file_objects: + self.logger.info(f"S3 metrics export verified via localstack: found {len(file_objects)} metric files") + + # Try to download and check the first file's content + try: + first_file_parts = file_objects[0].split() + if len(first_file_parts) >= 4: + object_name = ' '.join(first_file_parts[3:]) # File name may contain spaces + + # Download and check content + download_cmd = f"aws s3 cp s3://{bucket_name}/{object_name} /tmp/sample_metrics.json --endpoint=http://10.5.0.2:4566" + ret, download_output = subprocess.getstatusoutput(download_cmd) + if ret == 0: + self.logger.info(f"Successfully downloaded sample metrics file: {download_output}") + + # Check file content + cat_cmd = "head -n 3 /tmp/sample_metrics.json" + ret2, content = subprocess.getstatusoutput(cat_cmd) + if ret2 == 0: + self.logger.info(f"Sample metrics content: {content}") + # Verify content format is correct (should contain JSON formatted metric data) + if any(keyword in content for keyword in ['timestamp', 'name', 'kind', 'tags']): + self.logger.info("Metrics content format verification passed") + else: + self.logger.warning(f"Metrics content format may be incorrect: {content}") + else: + self.logger.warning(f"Failed to download sample file: {download_output}") + except Exception as e: + self.logger.warning(f"Error validating sample metrics file: {e}") + + return True + else: + self.logger.warning("Found S3 objects but none appear to be metric files") + return False + else: + # Check if bucket exists but is empty + bucket_check_cmd = f"aws s3api head-bucket --bucket {bucket_name} --endpoint-url http://10.5.0.2:4566" + ret, bucket_output = subprocess.getstatusoutput(bucket_check_cmd) + if ret == 0: + self.logger.info(f"Bucket {bucket_name} exists but is empty - metrics may not have been exported yet") + return False + else: + self.logger.warning(f"Bucket {bucket_name} may not exist: {bucket_output}") + return False + + except Exception as e: + self.logger.warning(f"Error verifying S3 metrics export via localstack: {e}") + return False + + @cluster(num_nodes=5) + def test_metrics_availability_basic(self): + """Test basic metrics endpoint availability""" + self.setup_services(num_workers=3) + self.cc.set_configs(lambda node: self.render("connect-distributed.properties", node=node)) + self.cc.start() + + try: + self.logger.info("Testing basic metrics availability...") + self._wait_for_metrics_available() + self.logger.info("Basic metrics availability test passed!") + + finally: + self.cc.stop() + + @cluster(num_nodes=5) + def test_opentelemetry_metrics_with_connector(self): + """Test OpenTelemetry metrics with running connector""" + self.setup_services(num_workers=3) + self.cc.set_configs(lambda node: self.render("connect-distributed.properties", node=node)) + self.cc.start() + + try: + self.logger.info("Creating VerifiableSource connector...") + # Use VerifiableSource connector + self.source = VerifiableSource(self.cc, topic='metrics-test-topic', throughput=10) + self.source.start() + + # Wait for connector to be running + self.logger.info("Waiting for connector to be running...") + wait_until(lambda: self.is_running(self.source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + self.logger.info("Connector is running, checking metrics...") + + # Wait for and verify metrics + self._wait_for_metrics_available() + self._verify_opentelemetry_metrics() + + # Verify metrics update over time + self._verify_metrics_updates() + + self.logger.info("All metrics validations passed!") + + finally: + if hasattr(self, 'source'): + self.logger.info("Stopping source connector...") + self.source.stop() + self.logger.info("Stopping Connect cluster...") + self.cc.stop() + + @cluster(num_nodes=5) + def test_comprehensive_metrics_validation(self): + """Comprehensive Connect OpenTelemetry metrics test""" + self.setup_services(num_workers=3) + self.cc.set_configs(lambda node: self.render("connect-distributed.properties", node=node)) + self.cc.start() + + try: + # Create connector using VerifiableSource + self.source = VerifiableSource(self.cc, topic='metrics-test-topic', throughput=50) + self.source.start() + + # Wait for connector startup + wait_until( + lambda: self.is_running(self.source), + timeout_sec=30, + err_msg="VerifiableSource connector failed to start within expected time" + ) + + # Verify metrics export + self._wait_for_metrics_available() + self._verify_comprehensive_metrics() + + # Verify connector is producing data + wait_until( + lambda: len(self.source.sent_messages()) > 0, + timeout_sec=30, + err_msg="VerifiableSource failed to produce messages" + ) + + finally: + if hasattr(self, 'source'): + self.source.stop() + self.cc.stop() + + def _connector_is_running_by_name(self, connector_name): + """Helper to check if connector is running by name""" + try: + status = self.cc.get_connector_status(connector_name) + return status and status['connector']['state'] == 'RUNNING' + except: + return False + + def _verify_metrics_under_load(self, expected_connector_count): + """Verify metrics accuracy under load""" + self._wait_for_metrics_available() + + for node in self.cc.nodes: + cmd = "curl -s http://localhost:9464/metrics" + result = node.account.ssh_capture(cmd) + metrics_output = "".join([line for line in result]) + + # Verify connector count metrics + connector_count_found = False + for line in metrics_output.split('\n'): + if 'kafka_connect_worker_connector_count' in line and not line.startswith('#'): + parts = line.split() + if len(parts) >= 2: + count = float(parts[1]) + assert count >= expected_connector_count, f"Connector count metric incorrect: {count} < {expected_connector_count}" + connector_count_found = True + break + + assert connector_count_found, "Connector count metric not found" + self.logger.info(f"Node {node.account.hostname} load test metrics validation passed") + + @cluster(num_nodes=5) + def test_opentelemetry_s3_metrics_exporter(self): + """Test OpenTelemetry S3 Metrics exporter functionality""" + # Setup mock S3 server using localstack + self.setup_services(num_workers=2) + cluster_id = f"connect-logs-{int(time.time())}" + bucket_name = "ko3" + metrics_prefix = f"automq/metrics/{cluster_id}" + + def s3_config(node): + config = self.render("connect-distributed.properties", node=node) + # Replace prometheus exporter with S3 exporter + config = config.replace( + "automq.telemetry.exporter.uri=prometheus://0.0.0.0:9464", + f"automq.telemetry.exporter.uri=ops://{bucket_name}" + ) + # Add S3 specific configurations + config += "\nautomq.telemetry.exporter.interval.ms=10000\n" + config += f"automq.telemetry.s3.bucket=0@s3://{bucket_name}?endpoint=http://10.5.0.2:4566®ion=us-east-1\n" + config += f"automq.telemetry.s3.cluster.id={cluster_id}\n" + config += f"automq.telemetry.s3.node.id={self.cc.nodes.index(node) + 1}\n" + config += "automq.telemetry.s3.selector.type=connect-leader\n" + + return config + + self.cc.set_configs(s3_config) + + def _list_s3_objects(prefix): + """List S3 objects with given prefix using kafka service method""" + objects, _ = self.kafka.get_bucket_objects() + return [obj for obj in objects if obj["path"].startswith(prefix)] + + try: + self.logger.info("Starting Connect cluster with S3 exporter...") + self.cc.start() + + def _connect_leader_nodes(): + leaders = [] + pattern = "Node became leader" + for connect_node in self.cc.nodes: + cmd = f"grep -a '{pattern}' {self.cc.LOG_FILE} || true" + output = "".join(connect_node.account.ssh_capture(cmd, allow_fail=True)) + if pattern in output: + leaders.append(connect_node.account.hostname) + return leaders + + wait_until( + lambda: len(_connect_leader_nodes()) == 1, + timeout_sec=120, + backoff_sec=5, + err_msg="Telemetry leadership in Connect cluster did not converge" + ) + + # Create connector to generate metrics + self.source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=15) + self.source.start() + + # Wait for connector to be running + wait_until(lambda: self.is_running(self.source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + # Wait for metrics to be exported to S3 + self.logger.info("Waiting for S3 metrics export...") + + def _metrics_uploaded(): + objects = _list_s3_objects(metrics_prefix) + if objects: + self.logger.info("Found %d metrics objects for prefix %s", len(objects), metrics_prefix) + return len(objects) > 0 + + wait_until( + _metrics_uploaded, + timeout_sec=180, + backoff_sec=10, + err_msg="Timed out waiting for Connect S3 metrics export" + ) + + self.logger.info("S3 Metrics exporter test passed!") + + finally: + # Cleanup + try: + if hasattr(self, 'source'): + self.source.stop() + self.cc.stop() + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + + @cluster(num_nodes=5) + def test_s3_log_uploader(self): + """Verify that Connect workers upload logs to S3 using the AutoMQ log uploader.""" + self.setup_services(num_workers=2) + + bucket_name = "ko3" + cluster_id = f"connect-logs-{int(time.time())}" + logs_prefix = f"automq/logs/{cluster_id}" + + def s3_log_config(node): + config = self.render("connect-distributed.properties", node=node) + config += "\nlog.s3.enable=true\n" + config += f"log.s3.bucket=0@s3://{bucket_name}?endpoint=http://10.5.0.2:4566®ion=us-east-1\n" + config += f"log.s3.cluster.id={cluster_id}\n" + config += f"log.s3.node.id={self.cc.nodes.index(node) + 1}\n" + config += "log.s3.selector.type=connect-leader\n" + + return config + + self.cc.set_configs(s3_log_config) + self.cc.environment['AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL'] = '15000' + self.cc.environment['AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL'] = '60000' + + def _list_s3_objects(prefix): + """List S3 objects with given prefix using kafka service method""" + objects, _ = self.kafka.get_bucket_objects() + return [obj for obj in objects if obj["path"].startswith(prefix)] + + source = None + + try: + self.logger.info("Starting Connect cluster with S3 log uploader enabled ...") + self.cc.start() + + def _connect_leader_nodes(): + leaders = [] + pattern = "Node became leader" + for connect_node in self.cc.nodes: + cmd = f"grep -a '{pattern}' {self.cc.LOG_FILE} || true" + output = "".join(connect_node.account.ssh_capture(cmd, allow_fail=True)) + if pattern in output: + leaders.append(connect_node.account.hostname) + return leaders + + wait_until( + lambda: len(_connect_leader_nodes()) == 1, + timeout_sec=120, + backoff_sec=5, + err_msg="Log uploader leadership in Connect cluster did not converge" + ) + + source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=10) + source.start() + + wait_until(lambda: self.is_running(source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + def _logs_uploaded(): + objects = _list_s3_objects(logs_prefix) + if objects: + self.logger.info("Found %d log objects for prefix %s", len(objects), logs_prefix) + return len(objects) > 0 + + wait_until(_logs_uploaded, timeout_sec=240, backoff_sec=15, + err_msg="Timed out waiting for Connect S3 log upload") + + # Verify objects are actually present + objects = _list_s3_objects(logs_prefix) + assert objects, "Expected log objects to be present after successful upload" + + finally: + try: + if source: + source.stop() + self.cc.stop() + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") diff --git a/tests/kafkatest/tests/connect/connect_remote_write_test.py b/tests/kafkatest/tests/connect/connect_remote_write_test.py new file mode 100644 index 0000000000..cedb3f9c65 --- /dev/null +++ b/tests/kafkatest/tests/connect/connect_remote_write_test.py @@ -0,0 +1,469 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ducktape.tests.test import Test +from ducktape.mark.resource import cluster +from ducktape.utils.util import wait_until + +from kafkatest.services.zookeeper import ZookeeperService +from kafkatest.services.kafka import KafkaService, quorum +from kafkatest.services.connect import ConnectDistributedService, VerifiableSource +from kafkatest.services.security.security_config import SecurityConfig +from kafkatest.version import DEV_BRANCH + +import time + + +class ConnectRemoteWriteTest(Test): + """ + Test cases for Kafka Connect OpenTelemetry Remote Write exporter functionality. + """ + + TOPIC = "remote-write-test" + FILE_SOURCE_CONNECTOR = 'org.apache.kafka.connect.file.FileStreamSourceConnector' + FILE_SINK_CONNECTOR = 'org.apache.kafka.connect.file.FileStreamSinkConnector' + + INPUT_FILE = "/mnt/connect.input" + OUTPUT_FILE = "/mnt/connect.output" + + TOPIC = "test" + OFFSETS_TOPIC = "connect-offsets" + OFFSETS_REPLICATION_FACTOR = "1" + OFFSETS_PARTITIONS = "1" + CONFIG_TOPIC = "connect-configs" + CONFIG_REPLICATION_FACTOR = "1" + STATUS_TOPIC = "connect-status" + STATUS_REPLICATION_FACTOR = "1" + STATUS_PARTITIONS = "1" + EXACTLY_ONCE_SOURCE_SUPPORT = "disabled" + SCHEDULED_REBALANCE_MAX_DELAY_MS = "60000" + CONNECT_PROTOCOL="sessioned" + + # Since tasks can be assigned to any node and we're testing with files, we need to make sure the content is the same + # across all nodes. + FIRST_INPUT_LIST = ["foo", "bar", "baz"] + FIRST_INPUTS = "\n".join(FIRST_INPUT_LIST) + "\n" + SECOND_INPUT_LIST = ["razz", "ma", "tazz"] + SECOND_INPUTS = "\n".join(SECOND_INPUT_LIST) + "\n" + + SCHEMA = { "type": "string", "optional": False } + + def __init__(self, test_context): + super(ConnectRemoteWriteTest, self).__init__(test_context) + self.num_zk = 1 + self.num_brokers = 1 + self.topics = { + self.TOPIC: {'partitions': 1, 'replication-factor': 1} + } + + self.zk = ZookeeperService(test_context, self.num_zk) if quorum.for_test(test_context) == quorum.zk else None + + def setup_services(self, num_workers=2): + self.kafka = KafkaService( + self.test_context, + self.num_brokers, + self.zk, + security_protocol=SecurityConfig.PLAINTEXT, + interbroker_security_protocol=SecurityConfig.PLAINTEXT, + topics=self.topics, + version=DEV_BRANCH, + server_prop_overrides=[ + ["auto.create.topics.enable", "false"], + ["transaction.state.log.replication.factor", str(self.num_brokers)], + ["transaction.state.log.min.isr", str(self.num_brokers)] + ], + allow_zk_with_kraft=True + ) + + self.cc = ConnectDistributedService( + self.test_context, + num_workers, + self.kafka, + ["/mnt/connect.input", "/mnt/connect.output"] + ) + self.cc.log_level = "DEBUG" + + if self.zk: + self.zk.start() + self.kafka.start() + + def is_running(self, connector, node=None): + """Check if connector is running""" + try: + status = self.cc.get_connector_status(connector.name, node) + return (status is not None and + status['connector']['state'] == 'RUNNING' and + all(task['state'] == 'RUNNING' for task in status['tasks'])) + except: + return False + + def _check_port_listening(self, node, port): + """Check if a port is listening on the given node, with multiple fallbacks""" + cmds = [ + f"ss -ltn | grep -E '(:|\\[::\\]):{port}\\b'", + f"netstat -ln | grep ':{port}\\b'", + f"lsof -iTCP:{port} -sTCP:LISTEN" + ] + for cmd in cmds: + try: + result = list(node.account.ssh_capture(cmd, allow_fail=True)) + if len(result) > 0: + return True + except Exception: + continue + return False + + def _start_mock_remote_write_server(self, node, port=9090, log_file="/tmp/mock_remote_write.log", + script_file="/tmp/mock_remote_write.py"): + """Start mock remote write HTTP server robustly""" + # Write script file (heredoc to avoid escaping issues) + write_cmd = f"""cat > {script_file} <<'PY' +import http.server +import socketserver +from urllib.parse import urlparse +import gzip +import sys +import time + +class MockRemoteWriteHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + if self.path == '/api/v1/write': + content_length = int(self.headers.get('Content-Length', 0)) + post_data = self.rfile.read(content_length) + encoding = self.headers.get('Content-Encoding', '') + if encoding == 'gzip': + try: + post_data = gzip.decompress(post_data) + except Exception: + pass + print(f"{{time.strftime('%Y-%m-%d-%H:%M:%S')}} - Received remote write request: {{len(post_data)}} bytes, encoding: {{encoding}}", flush=True) + self.send_response(200) + self.end_headers() + self.wfile.write(b'OK') + else: + print(f"{{time.strftime('%Y-%m-%d-%H:%M:%S')}} - Received non-write request: {{self.path}}", flush=True) + self.send_response(404) + self.end_headers() + + def log_message(self, format, *args): + print(f"{{time.strftime('%Y-%m-%d-%H:%M:%S')}} - HTTP: " + (format % args), flush=True) + +print('Mock remote write server starting...', flush=True) +with socketserver.TCPServer(('', {port}), MockRemoteWriteHandler) as httpd: + print('Mock remote write server listening on port {port}', flush=True) + httpd.serve_forever() +PY""" + node.account.ssh(write_cmd) + + # Choose python interpreter + which_py = "PYBIN=$(command -v python3 || command -v python || echo python3)" + # Start in background and record PID + start_cmd = f"{which_py}; nohup $PYBIN {script_file} > {log_file} 2>&1 & echo $!" + pid_out = list(node.account.ssh_capture(start_cmd)) + pid = pid_out[0].strip() if pid_out else None + if not pid: + raise RuntimeError("Failed to start mock remote write server (no PID)") + + # Wait for port to be listening + def listening(): + if self._check_port_listening(node, port): + return True + # If not listening, print recent logs for troubleshooting + try: + tail = "".join(list(node.account.ssh_capture(f"tail -n 20 {log_file}", allow_fail=True))) + self.logger.info(f"Mock server tail log: {tail}") + except Exception: + pass + return False + + wait_until(listening, timeout_sec=30, err_msg="Mock remote write server failed to start") + return pid + + def _verify_remote_write_requests(self, node, log_file="/tmp/mock_remote_write.log"): + """Verify that remote write requests were received""" + try: + # Check the mock server log for received requests + result = list(node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + + self.logger.info(f"Remote write log content: {log_content}") + + # Look for evidence of received data + if "Received" in log_content or "received" in log_content: + self.logger.info("Remote write requests were successfully received") + return True + + # Also check if the process is running and listening + if self._check_port_listening(node, 9090) or self._check_port_listening(node, 9091): + self.logger.info("Remote write server is listening, requests may have been processed") + return True + + self.logger.warning("No clear evidence of remote write requests in log") + return False + + except Exception as e: + self.logger.warning(f"Error verifying remote write requests: {e}") + # Don't fail the test if we can't verify the log, as the server might be working + return True + + @cluster(num_nodes=5) + def test_opentelemetry_remote_write_exporter(self): + """Test OpenTelemetry Remote Write exporter functionality""" + # Setup mock remote write server + self.setup_services(num_workers=2) + + # Override the template to use remote write exporter + def remote_write_config(node): + config = self.render("connect-distributed.properties", node=node) + # Replace prometheus exporter with remote write using correct URI format + self.logger.info(f"connect config: {config}") + config = config.replace( + "automq.telemetry.exporter.uri=prometheus://0.0.0.0:9464", + "automq.telemetry.exporter.uri=rw://?endpoint=http://localhost:9090/api/v1/write&auth=no_auth&maxBatchSize=1000000" + ) + # Add remote write specific configurations + config += "\nautomq.telemetry.exporter.interval.ms=30000\n" + + self.logger.info(f"connect new config: {config}") + return config + + self.cc.set_configs(remote_write_config) + + # Setup mock remote write endpoint + mock_server_node = self.cc.nodes[0] + self.logger.info("Setting up mock remote write server...") + + try: + # Start mock server + mock_pid = self._start_mock_remote_write_server(mock_server_node, port=9090) + self.logger.info(f"Mock remote write server started with PID: {mock_pid}") + + # Wait a bit for server to start + time.sleep(5) + + # Verify mock server is listening + wait_until( + lambda: self._check_port_listening(mock_server_node, 9090), + timeout_sec=30, + err_msg="Mock remote write server failed to start" + ) + + self.logger.info("Starting Connect cluster with Remote Write exporter...") + self.cc.start() + + # Create connector to generate metrics + self.source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=20) + self.source.start() + + # Wait for connector to be running + wait_until(lambda: self.is_running(self.source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + # Wait for metrics to be sent to remote write endpoint + self.logger.info("Waiting for remote write requests...") + time.sleep(120) # Wait for at least 2 export intervals + + # Verify remote write requests were received + self._verify_remote_write_requests(mock_server_node) + + self.logger.info("Remote Write exporter test passed!") + + finally: + # Cleanup + try: + if 'mock_pid' in locals() and mock_pid: + mock_server_node.account.ssh(f"kill {mock_pid}", allow_fail=True) + if hasattr(self, 'source'): + self.source.stop() + self.cc.stop() + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + + @cluster(num_nodes=5) + def test_remote_write_with_compression(self): + """Test remote write exporter with gzip compression""" + self.setup_services(num_workers=2) + + # Configure remote write with compression + def remote_write_config(node): + config = self.render("connect-distributed.properties", node=node) + config = config.replace( + "automq.telemetry.exporter.uri=prometheus://0.0.0.0:9464", + "automq.telemetry.exporter.uri=rw://?endpoint=http://localhost:9091/api/v1/write&auth=no_auth&maxBatchSize=500000&compression=gzip" + ) + config += "\nautomq.telemetry.exporter.interval.ms=20000\n" + return config + + self.cc.set_configs(remote_write_config) + + mock_server_node = self.cc.nodes[0] + + try: + # Start mock server on different port + mock_pid = self._start_mock_remote_write_server(mock_server_node, port=9091) + + wait_until( + lambda: self._check_port_listening(mock_server_node, 9091), + timeout_sec=30, + err_msg="Mock remote write server failed to start" + ) + + self.cc.start() + + # Create connector + self.source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=30) + self.source.start() + + wait_until(lambda: self.is_running(self.source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + # Wait for compressed requests + time.sleep(100) + + # Verify requests were received + log_file = "/tmp/mock_remote_write.log" + assert self._verify_remote_write_requests(mock_server_node, log_file), \ + "Did not observe remote write payloads at the mock endpoint" + + # Check for gzip compression evidence + result = list(mock_server_node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + if "encoding: gzip" in log_content: + self.logger.info("Verified gzip compression was used for remote write requests") + else: + self.logger.warning("No evidence of gzip compression in remote write requests") + + self.logger.info("Remote write compression test passed!") + + finally: + try: + if 'mock_pid' in locals() and mock_pid: + mock_server_node.account.ssh(f"kill {mock_pid}", allow_fail=True) + if hasattr(self, 'source'): + self.source.stop() + self.cc.stop() + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + + @cluster(num_nodes=5) + def test_remote_write_batch_size_limits(self): + """Test remote write exporter with different batch size configurations""" + self.setup_services(num_workers=2) + + # Test with smaller batch size to ensure multiple requests + def remote_write_config(node): + config = self.render("connect-distributed.properties", node=node) + config = config.replace( + "automq.telemetry.exporter.uri=prometheus://0.0.0.0:9464", + "automq.telemetry.exporter.uri=rw://?endpoint=http://localhost:9092/api/v1/write&auth=no_auth&maxBatchSize=10000" + ) + config += "\nautomq.telemetry.exporter.interval.ms=15000\n" + return config + + self.cc.set_configs(remote_write_config) + + mock_server_node = self.cc.nodes[0] + + try: + mock_pid = self._start_mock_remote_write_server(mock_server_node, port=9092) + + wait_until( + lambda: self._check_port_listening(mock_server_node, 9092), + timeout_sec=30, + err_msg="Mock remote write server failed to start" + ) + + self.cc.start() + + # Create connector with higher throughput to generate more metrics + self.source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=100) + self.source.start() + + wait_until(lambda: self.is_running(self.source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + # Wait for multiple batched requests + time.sleep(90) + + # Verify multiple requests were received due to batch size limits + log_file = "/tmp/mock_remote_write.log" + result = list(mock_server_node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + + # Count the number of received requests + request_count = log_content.count("Received remote write request") + self.logger.info(f"Received {request_count} remote write requests") + + assert request_count > 1, f"Expected multiple remote write requests due to batch size limits, but only received {request_count}" + + self.logger.info("Remote write batch size test passed!") + + finally: + try: + if 'mock_pid' in locals() and mock_pid: + mock_server_node.account.ssh(f"kill {mock_pid}", allow_fail=True) + if hasattr(self, 'source'): + self.source.stop() + self.cc.stop() + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + + @cluster(num_nodes=5) + def test_remote_write_server_unavailable(self): + """Test remote write exporter behavior when server is unavailable""" + self.setup_services(num_workers=2) + + # Configure remote write to point to unavailable server + def remote_write_config(node): + config = self.render("connect-distributed.properties", node=node) + config = config.replace( + "automq.telemetry.exporter.uri=prometheus://0.0.0.0:9464", + "automq.telemetry.exporter.uri=rw://?endpoint=http://localhost:9999/api/v1/write&auth=no_auth&maxBatchSize=1000000" + ) + config += "\nautomq.telemetry.exporter.interval.ms=10000\n" + return config + + self.cc.set_configs(remote_write_config) + + try: + self.logger.info("Testing remote write behavior with unavailable server...") + self.cc.start() + + # Create connector even though remote write server is unavailable + self.source = VerifiableSource(self.cc, topic=self.TOPIC, throughput=20) + self.source.start() + + wait_until(lambda: self.is_running(self.source), timeout_sec=30, + err_msg="VerifiableSource connector failed to start") + + # Wait for export attempts + time.sleep(60) + + # Kafka Connect should continue functioning normally even if remote write fails + # This is primarily a resilience test - we verify the connector doesn't crash + self.logger.info("Connector remained stable with unavailable remote write server") + + # Verify connector is still responsive + assert self.is_running(self.source), "Connector should remain running despite remote write failures" + + self.logger.info("Remote write unavailable server test passed!") + + finally: + try: + if hasattr(self, 'source'): + self.source.stop() + self.cc.stop() + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") diff --git a/tests/kafkatest/tests/connect/templates/connect-distributed.properties b/tests/kafkatest/tests/connect/templates/connect-distributed.properties index fa2172edd7..724c6a971a 100644 --- a/tests/kafkatest/tests/connect/templates/connect-distributed.properties +++ b/tests/kafkatest/tests/connect/templates/connect-distributed.properties @@ -69,4 +69,13 @@ config.providers.file.class=org.apache.kafka.common.config.provider.FileConfigPr {% if PLUGIN_PATH is defined %} plugin.path={{ PLUGIN_PATH }} {% endif %} -plugin.discovery={{ PLUGIN_DISCOVERY|default("service_load") }} \ No newline at end of file +plugin.discovery={{ PLUGIN_DISCOVERY|default("service_load") }} + +metric.reporters=org.apache.kafka.connect.automq.metrics.OpenTelemetryMetricsReporter + +opentelemetry.metrics.enabled=true +opentelemetry.metrics.prefix=kafka.connect + +automq.telemetry.exporter.uri=prometheus://0.0.0.0:9464 +service.name=kafka-connect-test +service.instance.id=1 \ No newline at end of file diff --git a/tests/kafkatest/tests/core/automq_remote_write_test.py b/tests/kafkatest/tests/core/automq_remote_write_test.py new file mode 100644 index 0000000000..1522373d13 --- /dev/null +++ b/tests/kafkatest/tests/core/automq_remote_write_test.py @@ -0,0 +1,386 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import time + +from ducktape.mark.resource import cluster +from ducktape.tests.test import Test +from ducktape.utils.util import wait_until + +from kafkatest.services.kafka import KafkaService, quorum +from kafkatest.services.security.security_config import SecurityConfig +from kafkatest.services.verifiable_producer import VerifiableProducer +from kafkatest.services.zookeeper import ZookeeperService + + +class AutoMQRemoteWriteTest(Test): + """End-to-end validation for AutoMQ Remote Write exporter integration.""" + + TOPIC = "automq-remote-write-topic" + + def __init__(self, test_context): + super(AutoMQRemoteWriteTest, self).__init__(test_context) + self.num_brokers = 1 + self.zk = None + self.kafka = None + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _start_kafka(self, server_overrides=None, per_node_overrides=None, extra_env=None): + if quorum.for_test(self.test_context) == quorum.zk and self.zk is None: + self.zk = ZookeeperService(self.test_context, 1) + self.zk.start() + + self.kafka = KafkaService( + self.test_context, + self.num_brokers, + self.zk, + security_protocol=SecurityConfig.PLAINTEXT, + topics={}, + server_prop_overrides=server_overrides, + per_node_server_prop_overrides=per_node_overrides, + extra_env=extra_env, + ) + + self.kafka.start() + self.kafka.create_topic({ + "topic": self.TOPIC, + "partitions": 1, + "replication-factor": 1, + }) + + def _stop_kafka(self): + if self.kafka is not None: + self.kafka.stop() + self.kafka = None + if self.zk is not None: + self.zk.stop() + self.zk = None + + def _produce_messages(self, max_messages=200, throughput=1000): + producer = VerifiableProducer( + self.test_context, + num_nodes=1, + kafka=self.kafka, + topic=self.TOPIC, + max_messages=max_messages, + throughput=throughput, + ) + producer.start() + try: + wait_until( + lambda: producer.num_acked >= max_messages, + timeout_sec=60, + backoff_sec=5, + err_msg="Producer failed to deliver expected number of messages", + ) + finally: + try: + producer.stop() + except Exception as e: + self.logger.warn("Error stopping producer: %s", e) + + def _check_port_listening(self, node, port): + """Check if a port is listening on the given node, with multiple fallbacks""" + cmds = [ + f"ss -ltn | grep -E '(:|\\[::\\]):{port}\\b'", + f"netstat -ln | grep ':{port}\\b'", + f"lsof -iTCP:{port} -sTCP:LISTEN" + ] + for cmd in cmds: + try: + result = list(node.account.ssh_capture(cmd, allow_fail=True)) + if len(result) > 0: + return True + except Exception: + continue + return False + + def _verify_remote_write_requests(self, node, log_file): + """Verify that remote write requests were captured by the mock server.""" + try: + result = list(node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + if "Received" in log_content: + self.logger.info("Remote write server captured payloads: %s", log_content) + return True + self.logger.warning("No remote write payload entries detected in %s", log_file) + return False + except Exception as e: + self.logger.warning("Failed to read remote write log %s: %s", log_file, e) + return False + + def _start_mock_remote_write_server(self, node, port=9090, log_file="/tmp/mock_remote_write.log", + script_file="/tmp/mock_remote_write.py"): + """Start mock remote write HTTP server robustly""" + # Write script file (heredoc to avoid escaping issues) + write_cmd = f"""cat > {script_file} <<'PY' +import http.server +import socketserver +from urllib.parse import urlparse +import gzip +import sys +import time + +class MockRemoteWriteHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + if self.path == '/api/v1/write': + content_length = int(self.headers.get('Content-Length', 0)) + post_data = self.rfile.read(content_length) + encoding = self.headers.get('Content-Encoding', '') + if encoding == 'gzip': + try: + post_data = gzip.decompress(post_data) + except Exception: + pass + print(f"{{time.strftime('%Y-%m-%d-%H:%M:%S')}} - Received remote write request: {{len(post_data)}} bytes, encoding: {{encoding}}", flush=True) + self.send_response(200) + self.end_headers() + self.wfile.write(b'OK') + else: + print(f"{{time.strftime('%Y-%m-%d-%H:%M:%S')}} - Received non-write request: {{self.path}}", flush=True) + self.send_response(404) + self.end_headers() + + def log_message(self, format, *args): + print(f"{{time.strftime('%Y-%m-%d-%H:%M:%S')}} - HTTP: " + (format % args), flush=True) + +print('Mock remote write server starting...', flush=True) +with socketserver.TCPServer(('', {port}), MockRemoteWriteHandler) as httpd: + print('Mock remote write server listening on port {port}', flush=True) + httpd.serve_forever() +PY""" + node.account.ssh(write_cmd) + + # Choose python interpreter + which_py = "PYBIN=$(command -v python3 || command -v python || echo python3)" + # Start in background and record PID + start_cmd = f"{which_py}; nohup $PYBIN {script_file} > {log_file} 2>&1 & echo $!" + pid_out = list(node.account.ssh_capture(start_cmd)) + pid = pid_out[0].strip() if pid_out else None + if not pid: + raise RuntimeError("Failed to start mock remote write server (no PID)") + + # Wait for port to be listening + def listening(): + if self._check_port_listening(node, port): + return True + # If not listening, print recent logs for troubleshooting + try: + tail = "".join(list(node.account.ssh_capture(f"tail -n 20 {log_file}", allow_fail=True))) + self.logger.info(f"Mock server tail log: {tail}") + except Exception: + pass + return False + + wait_until(listening, timeout_sec=30, err_msg="Mock remote write server failed to start") + return pid + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + @cluster(num_nodes=5) + def test_remote_write_metrics_exporter(self): + """Verify remote write exporter integration using a mock HTTP endpoint.""" + cluster_id = f"core-remote-write-{int(time.time())}" + remote_write_port = 19090 + log_file = f"/tmp/automq_remote_write_{int(time.time())}.log" + script_path = f"/tmp/automq_remote_write_server_{int(time.time())}.py" + + server_overrides = [ + ["s3.telemetry.metrics.exporter.uri", f"rw://?endpoint=http://localhost:{remote_write_port}/api/v1/write&auth=no_auth&maxBatchSize=1000000"], + ["s3.telemetry.exporter.report.interval.ms", "15000"], + ["service.name", cluster_id], + ["service.instance.id", "broker-remote-write"], + ] + + remote_write_node = None + mock_pid = None + + self._start_kafka(server_overrides=server_overrides) + + try: + remote_write_node = self.kafka.nodes[0] + self.logger.info("Setting up mock remote write server...") + + mock_pid = self._start_mock_remote_write_server(remote_write_node, remote_write_port, log_file, script_path) + + self.logger.info("Starting message production...") + self._produce_messages(max_messages=400, throughput=800) + + # Allow multiple export intervals + self.logger.info("Waiting for remote write requests...") + time.sleep(120) + + assert self._verify_remote_write_requests(remote_write_node, log_file), \ + "Did not observe remote write payloads at the mock endpoint" + + self.logger.info("Remote write exporter test passed!") + finally: + try: + if remote_write_node is not None and mock_pid: + remote_write_node.account.ssh(f"kill {mock_pid}", allow_fail=True) + if remote_write_node is not None: + remote_write_node.account.ssh(f"rm -f {script_path}", allow_fail=True) + remote_write_node.account.ssh(f"rm -f {log_file}", allow_fail=True) + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + self._stop_kafka() + + @cluster(num_nodes=5) + def test_remote_write_with_compression(self): + """Test remote write exporter with gzip compression enabled.""" + cluster_id = f"core-remote-write-gzip-{int(time.time())}" + remote_write_port = 19091 + log_file = f"/tmp/automq_remote_write_gzip_{int(time.time())}.log" + script_path = f"/tmp/automq_remote_write_gzip_server_{int(time.time())}.py" + + server_overrides = [ + ["s3.telemetry.metrics.exporter.uri", f"rw://?endpoint=http://localhost:{remote_write_port}/api/v1/write&auth=no_auth&maxBatchSize=500000&compression=gzip"], + ["s3.telemetry.exporter.report.interval.ms", "10000"], + ["service.name", cluster_id], + ["service.instance.id", "broker-remote-write-gzip"], + ] + + self._start_kafka(server_overrides=server_overrides) + + try: + remote_write_node = self.kafka.nodes[0] + self.logger.info("Setting up mock remote write server with compression support...") + + mock_pid = self._start_mock_remote_write_server(remote_write_node, remote_write_port, log_file, script_path) + + self.logger.info("Starting message production for compression test...") + self._produce_messages(max_messages=600, throughput=1000) + + self.logger.info("Waiting for compressed remote write requests...") + time.sleep(90) + + # Verify requests were received + assert self._verify_remote_write_requests(remote_write_node, log_file), \ + "Did not observe compressed remote write payloads at the mock endpoint" + + # Check that gzip encoding was used + result = list(remote_write_node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + if "encoding: gzip" in log_content: + self.logger.info("Verified gzip compression was used for remote write requests") + else: + self.logger.warning("No evidence of gzip compression in remote write requests") + + self.logger.info("Remote write compression test passed!") + finally: + try: + if 'remote_write_node' in locals() and 'mock_pid' in locals() and mock_pid: + remote_write_node.account.ssh(f"kill {mock_pid}", allow_fail=True) + if 'remote_write_node' in locals(): + remote_write_node.account.ssh(f"rm -f {script_path}", allow_fail=True) + remote_write_node.account.ssh(f"rm -f {log_file}", allow_fail=True) + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + self._stop_kafka() + + @cluster(num_nodes=5) + def test_remote_write_batch_size_limits(self): + """Test remote write exporter with different batch size configurations.""" + cluster_id = f"core-remote-write-batch-{int(time.time())}" + remote_write_port = 19092 + log_file = f"/tmp/automq_remote_write_batch_{int(time.time())}.log" + script_path = f"/tmp/automq_remote_write_batch_server_{int(time.time())}.py" + + # Test with smaller batch size to ensure multiple requests + server_overrides = [ + ["s3.telemetry.metrics.exporter.uri", f"rw://?endpoint=http://localhost:{remote_write_port}/api/v1/write&auth=no_auth&maxBatchSize=10000"], + ["s3.telemetry.exporter.report.interval.ms", "5000"], + ["service.name", cluster_id], + ["service.instance.id", "broker-remote-write-batch"], + ] + + self._start_kafka(server_overrides=server_overrides) + + try: + remote_write_node = self.kafka.nodes[0] + self.logger.info("Setting up mock remote write server for batch size testing...") + + mock_pid = self._start_mock_remote_write_server(remote_write_node, remote_write_port, log_file, script_path) + + self.logger.info("Starting high-volume message production...") + # Produce more messages to trigger multiple batches + self._produce_messages(max_messages=1000, throughput=2000) + + self.logger.info("Waiting for multiple batched remote write requests...") + time.sleep(60) + + # Verify multiple requests were received due to batch size limits + result = list(remote_write_node.account.ssh_capture(f"cat {log_file}", allow_fail=True)) + log_content = "".join(result) + + # Count the number of received requests + request_count = log_content.count("Received remote write request") + self.logger.info(f"Received {request_count} remote write requests") + + assert request_count > 1, f"Expected multiple remote write requests due to batch size limits, but only received {request_count}" + + self.logger.info("Remote write batch size test passed!") + finally: + try: + if 'remote_write_node' in locals() and 'mock_pid' in locals() and mock_pid: + remote_write_node.account.ssh(f"kill {mock_pid}", allow_fail=True) + if 'remote_write_node' in locals(): + remote_write_node.account.ssh(f"rm -f {script_path}", allow_fail=True) + remote_write_node.account.ssh(f"rm -f {log_file}", allow_fail=True) + except Exception as e: + self.logger.warning(f"Cleanup error: {e}") + self._stop_kafka() + + @cluster(num_nodes=5) + def test_remote_write_server_unavailable(self): + """Test remote write exporter behavior when server is unavailable.""" + cluster_id = f"core-remote-write-unavail-{int(time.time())}" + # Use a port that we won't start a server on + remote_write_port = 19093 + + server_overrides = [ + ["s3.telemetry.metrics.exporter.uri", f"rw://?endpoint=http://localhost:{remote_write_port}/api/v1/write&auth=no_auth&maxBatchSize=1000000"], + ["s3.telemetry.exporter.report.interval.ms", "10000"], + ["service.name", cluster_id], + ["service.instance.id", "broker-remote-write-unavail"], + ] + + self._start_kafka(server_overrides=server_overrides) + + try: + self.logger.info("Testing remote write behavior with unavailable server...") + + # Produce messages even though remote write server is unavailable + self._produce_messages(max_messages=200, throughput=500) + + # Wait for export attempts + time.sleep(30) + + # Kafka should continue functioning normally even if remote write fails + # This is primarily a resilience test - we verify the broker doesn't crash + self.logger.info("Broker remained stable with unavailable remote write server") + + # Verify broker is still responsive + final_messages = 100 + self._produce_messages(max_messages=final_messages, throughput=200) + + self.logger.info("Remote write unavailable server test passed!") + finally: + self._stop_kafka() diff --git a/tests/kafkatest/tests/core/automq_tabletopic_base.py b/tests/kafkatest/tests/core/automq_tabletopic_base.py new file mode 100644 index 0000000000..acce98ad27 --- /dev/null +++ b/tests/kafkatest/tests/core/automq_tabletopic_base.py @@ -0,0 +1,292 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import time +from typing import Dict, Any + +import requests +from ducktape.tests.test import Test +from ducktape.utils.util import wait_until + +from kafkatest.services.external_services import DockerComposeService +from kafkatest.services.kafka import KafkaService +from kafkatest.services.performance import AutoMQPerformanceService +from kafkatest.services.security.security_config import SecurityConfig + + +# Static IP configuration from ducknet +SCHEMA_REGISTRY_HOST = "10.5.1.3" +SCHEMA_REGISTRY_PORT = 8081 +SCHEMA_REGISTRY_URL = f"http://{SCHEMA_REGISTRY_HOST}:{SCHEMA_REGISTRY_PORT}" + +ICEBERG_CATALOG_HOST = "10.5.1.4" +ICEBERG_CATALOG_PORT = 8181 +ICEBERG_CATALOG_URL = f"http://{ICEBERG_CATALOG_HOST}:{ICEBERG_CATALOG_PORT}" +ICEBERG_CATALOG_S3_BUCKET = "ko3" +ICEBERG_CATALOG_DB_NAME = "default" + +# Docker compose file paths +DOCKER_COMPOSE_BASE_PATH = "/opt/kafka-dev/tests/docker" +SCHEMA_REGISTRY_COMPOSE_PATH = os.path.join(DOCKER_COMPOSE_BASE_PATH, "schema-registry/docker-compose.yaml") +ICEBERG_CATALOG_COMPOSE_PATH = os.path.join(DOCKER_COMPOSE_BASE_PATH, "iceberg-catalog/docker-compose.yaml") + + +class TableTopicBase(Test): + """Common base for TableTopic tests: handles service lifecycle and shared helpers.""" + + def __init__(self, test_context): + super(TableTopicBase, self).__init__(test_context) + self.iceberg_catalog_service = None + self.schema_registry_service = None + + self.security_protocol = SecurityConfig.PLAINTEXT + self.interbroker_security_protocol = SecurityConfig.PLAINTEXT + + self.kafka = KafkaService( + test_context, + num_nodes=3, + zk=None, + security_protocol=self.security_protocol, + interbroker_security_protocol=self.interbroker_security_protocol, + server_prop_overrides=[ + ["automq.table.topic.catalog.type", "rest"], + ["automq.table.topic.catalog.uri", ICEBERG_CATALOG_URL], + ["automq.table.topic.catalog.warehouse", f"s3://{ICEBERG_CATALOG_S3_BUCKET}/wh/"], + ["automq.table.topic.schema.registry.url", SCHEMA_REGISTRY_URL] + ], + topics={} + ) + + def setUp(self): + self.iceberg_catalog_service = DockerComposeService(ICEBERG_CATALOG_COMPOSE_PATH, self.logger) + self.iceberg_catalog_service.start() + + # Apply any security overrides configured by subclasses prior to starting Kafka + self.configure_security(self.security_protocol, self.interbroker_security_protocol, restart_cluster=False) + self.kafka.start() + bootstrap_servers = self.kafka.bootstrap_servers() + schema_registry_env = {"KAFKA_BOOTSTRAP_SERVERS": bootstrap_servers} + self.schema_registry_service = DockerComposeService(SCHEMA_REGISTRY_COMPOSE_PATH, self.logger) + self.schema_registry_service.start(env=schema_registry_env) + + self._wait_for_service(ICEBERG_CATALOG_URL + "/v1/config", "Iceberg Catalog", 200) + self._wait_for_service(SCHEMA_REGISTRY_URL, "Schema Registry") + + def tearDown(self): + if self.schema_registry_service: + self.schema_registry_service.stop() + if self.iceberg_catalog_service: + self.iceberg_catalog_service.stop() + + def _wait_for_service(self, url, service_name, expected_status=200, timeout_sec=60): + self.logger.info(f"Waiting for {service_name} at {url}...") + wait_until(lambda: self._ok(url, expected_status), timeout_sec=timeout_sec, + err_msg=f"{service_name} not ready in {timeout_sec}s") + + def _ok(self, url, expected_status): + try: + r = requests.get(url, timeout=5) + return r.status_code == expected_status + except Exception: + return False + + def _perf_topic_name(self, prefix: str, partitions: int = 16, index: int = 0) -> str: + return f"__automq_perf_{prefix}_{partitions:04d}_{index:07d}" + + def configure_security(self, security_protocol=None, interbroker_security_protocol=None, restart_cluster=False): + changed = False + if security_protocol is not None and security_protocol != self.kafka.security_protocol: + self.kafka.security_protocol = security_protocol + self.kafka.open_port(self.kafka.security_protocol) + changed = True + if interbroker_security_protocol is not None and interbroker_security_protocol != self.kafka.interbroker_security_protocol: + use_separate = interbroker_security_protocol != self.kafka.security_protocol + self.kafka.setup_interbroker_listener(interbroker_security_protocol, use_separate_listener=use_separate) + self.kafka.open_port(self.kafka.interbroker_listener.name) + self.kafka.interbroker_listener.open = True + changed = True + if changed: + self.kafka._security_config = None + if restart_cluster and changed: + if any(self.kafka.alive(node) for node in self.kafka.nodes): + self.kafka.restart_cluster() + else: + self.logger.info("Kafka not started; skipping restart after security change") + return changed + + def _sum_kafka_end_offsets(self, topic: str) -> int: + output = self.kafka.get_offset_shell(time='-1', topic=topic) + total = 0 + for raw in output.splitlines(): + line = raw.strip() + parts = line.split(":") + if len(parts) >= 3 and parts[0] == topic: + try: + total += int(parts[-1]) + except Exception: + pass + return total + + def _sum_kafka_end_offsets_stable(self, topic: str, settle_window_sec: int = 2, timeout_sec: int = 15) -> int: + start = time.time() + last = self._sum_kafka_end_offsets(topic) + last_change = time.time() + while time.time() - start < timeout_sec: + time.sleep(0.5) + cur = self._sum_kafka_end_offsets(topic) + if cur != last: + last = cur + last_change = time.time() + else: + if time.time() - last_change >= settle_window_sec: + return cur + return last + + def _verify_data_in_iceberg(self, topic_name: str, expected_count: int, commit_interval_ms: int = 2000): + """Verify Iceberg table exists and poll until snapshot shows expected_count or timeout.""" + table_url = f"{ICEBERG_CATALOG_URL}/v1/namespaces/{ICEBERG_CATALOG_DB_NAME}/tables/{topic_name}" + snapshots_url = f"{table_url}/snapshots" + headers = {'Accept': 'application/vnd.iceberg.v1+json, application/json'} + + timeout_sec = max(30, int(3 * commit_interval_ms / 1000)) + deadline = time.time() + timeout_sec + last_record_count = None + last_status = None + last_error = None + last_snapshot_id = None + last_num_snapshots = None + last_metadata_location = None + + def parse(meta: Dict[str, Any]): + info = {'record_count': None, 'snapshot_id': None, 'num_snapshots': None, 'metadata_location': None} + if not isinstance(meta, dict): + return info + m = meta.get('metadata') or meta + info['metadata_location'] = m.get('metadata-location') or m.get('metadataLocation') or \ + meta.get('metadata-location') or meta.get('metadataLocation') + cur = m.get('current-snapshot') or m.get('currentSnapshot') + if isinstance(cur, dict): + s = cur.get('summary') or {} + total = s.get('total-records') or s.get('totalRecords') + try: + info['record_count'] = int(total) if total is not None else None + except Exception: + info['record_count'] = None + info['snapshot_id'] = cur.get('snapshot-id') or cur.get('snapshotId') + return info + cur_id = m.get('current-snapshot-id') or m.get('currentSnapshotId') + snaps = m.get('snapshots') or [] + info['num_snapshots'] = len(snaps) if isinstance(snaps, list) else None + if cur_id and isinstance(snaps, list): + for s in snaps: + sid = s.get('snapshot-id') or s.get('snapshotId') + if sid == cur_id: + summ = s.get('summary') or {} + total = summ.get('total-records') or summ.get('totalRecords') + try: + info['record_count'] = int(total) if total is not None else None + except Exception: + info['record_count'] = None + info['snapshot_id'] = cur_id + return info + if isinstance(snaps, list) and snaps: + s = snaps[-1] + summ = s.get('summary') or {} + total = summ.get('total-records') or summ.get('totalRecords') + try: + info['record_count'] = int(total) if total is not None else None + except Exception: + info['record_count'] = None + info['snapshot_id'] = s.get('snapshot-id') or s.get('snapshotId') + return info + + while time.time() < deadline: + try: + resp = requests.get(table_url, headers=headers, timeout=5) + last_status = resp.status_code + if resp.status_code != 200: + time.sleep(1) + continue + meta = resp.json() + p = parse(meta) + last_record_count = p['record_count'] + last_snapshot_id = p['snapshot_id'] + last_num_snapshots = p['num_snapshots'] + last_metadata_location = p['metadata_location'] + self.logger.info( + f"Iceberg table ok. metaLocation={last_metadata_location}, snapshots={last_num_snapshots}, " + f"currentSnapshot={last_snapshot_id}, total-records={last_record_count} (expect {expected_count})") + if last_record_count == expected_count: + break + # optional snapshots endpoint + if last_num_snapshots is None: + try: + sresp = requests.get(snapshots_url, headers=headers, timeout=5) + self.logger.info(f"Iceberg GET {snapshots_url} -> {sresp.status_code}") + except Exception: + pass + time.sleep(1) + except requests.exceptions.RequestException as e: + last_error = str(e) + time.sleep(1) + + if last_record_count != expected_count: + msg = ( + f"Expected {expected_count} records, but found {last_record_count} in Iceberg after {timeout_sec}s. " + f"status={last_status}, snapshotId={last_snapshot_id}, snapshots={last_num_snapshots}, " + f"metaLocation={last_metadata_location}, lastError={last_error}. URL={table_url}" + ) + self.logger.error(msg) + assert False, msg + + def default_tabletopic_configs(self, commit_interval_ms: int, namespace: str = ICEBERG_CATALOG_DB_NAME) -> Dict[str, str]: + return { + "automq.table.topic.enable": "true", + "automq.table.topic.commit.interval.ms": str(commit_interval_ms), + "automq.table.topic.convert.value.type": "by_schema_id", + "automq.table.topic.transform.value.type": "flatten", + "automq.table.topic.namespace": namespace, + } + + def _run_perf(self, topic_prefix: str, topics: int, partitions: int, send_rate: int, + commit_interval_ms: int, duration_seconds: int = None, + value_schema_json: str = None, topic_configs: Dict[str, str] = None): + if duration_seconds is None: + duration_seconds = 60 + perf = AutoMQPerformanceService( + self.test_context, 1, kafka=self.kafka, + producers_per_topic=1, groups_per_topic=0, consumers_per_group=1, + topics=topics, partitions_per_topic=partitions, + send_rate=send_rate, record_size=256, + topic_prefix=topic_prefix, await_topic_ready=False, + topic_configs=(topic_configs or {}), + producer_configs={ + "schema.registry.url": SCHEMA_REGISTRY_URL, + "auto.register.schemas": "true", + }, + consumer_configs={}, + warmup_duration_minutes=0, + test_duration_minutes=max(1, int((duration_seconds + 59) // 60)), + value_schema=value_schema_json, + ) + try: + perf.run() + finally: + try: + perf.stop() + except Exception: + pass diff --git a/tests/kafkatest/tests/core/automq_tabletopic_broker_restart_test.py b/tests/kafkatest/tests/core/automq_tabletopic_broker_restart_test.py new file mode 100644 index 0000000000..9f863a2d6e --- /dev/null +++ b/tests/kafkatest/tests/core/automq_tabletopic_broker_restart_test.py @@ -0,0 +1,83 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import random +import threading +import time + +from .automq_tabletopic_base import TableTopicBase + + +class TableTopicBrokerRestartTest(TableTopicBase): + def __init__(self, test_context): + super(TableTopicBrokerRestartTest, self).__init__(test_context) + self.avro_schema = { + "type": "record", + "name": "RestartAvroSchema", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "ts", "type": "long"} + ] + } + + def _restart_brokers_periodically(self, interval_secs: int, total_secs: int): + stop_event = threading.Event() + + def loop(): + start = time.time() + while not stop_event.is_set() and time.time() - start < total_secs: + time.sleep(interval_secs) + try: + node = random.choice(self.kafka.nodes) + self.logger.info(f"[PeriodicRestart] Restart broker {node.account.hostname}") + # Clean restart + self.kafka.stop_node(node, clean_shutdown=True, timeout_sec=60) + self.kafka.start_node(node, timeout_sec=60) + except Exception as e: + self.logger.warn(f"[PeriodicRestart] failed: {e}") + + th = threading.Thread(target=loop, daemon=True) + th.start() + return stop_event, th + + def test_broker_periodic_restart_10m(self): + topic_prefix = "tt_restart" + topics = 1 + partitions = 16 + send_rate = 200 + commit_interval_ms = 2000 + test_duration_seconds = 120 # 10 minutes + + # Start periodic restarts every 30s + stop_event, th = self._restart_brokers_periodically(interval_secs=30, total_secs=test_duration_seconds) + + topic_configs = self.default_tabletopic_configs(commit_interval_ms) + + try: + # Run perf using base helper (blocks for duration) + self._run_perf(topic_prefix, topics, partitions, send_rate, + commit_interval_ms, test_duration_seconds, + value_schema_json=json.dumps(self.avro_schema), + topic_configs=topic_configs) + finally: + # Stop periodic restarts + stop_event.set() + th.join(timeout=10) + + topic_name = self._perf_topic_name(topic_prefix, partitions, 0) + produced = int(self._sum_kafka_end_offsets_stable(topic_name, 3, 60)) + self._verify_data_in_iceberg(topic_name, produced, commit_interval_ms=commit_interval_ms) diff --git a/tests/kafkatest/tests/core/automq_tabletopic_e2e_test.py b/tests/kafkatest/tests/core/automq_tabletopic_e2e_test.py new file mode 100644 index 0000000000..16ce23b68a --- /dev/null +++ b/tests/kafkatest/tests/core/automq_tabletopic_e2e_test.py @@ -0,0 +1,81 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from ducktape.mark import parametrize +from .automq_tabletopic_base import TableTopicBase +from kafkatest.services.security.security_config import SecurityConfig +from kafkatest.version import DEV_BRANCH + + +class TableTopicE2ETest(TableTopicBase): + """ + End-to-end test for the Table Topic feature. + Manages service lifecycle manually in setUp and tearDown to ensure correct startup order. + """ + def __init__(self, test_context): + super(TableTopicE2ETest, self).__init__(test_context) + + self.base_topic = "test_table_topic" + + # Avro schema definition + self.avro_schema = { + "type": "record", + "name": "TestAvroSchema", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "timestamp", "type": "long"} + ] + } + + def _produce_avro_data_with_perf(self, topic_prefix: str, value_schema_json: str, target_msgs: int = 50, + send_rate: int = 20, test_duration_seconds: int = None, + commit_interval_ms: int = 2000) -> (str, int): + if test_duration_seconds is None: + test_duration_seconds = max(60, int((target_msgs + send_rate - 1) // send_rate)) + # Use base helper to run perf with standard TableTopic configs + topic_configs = self.default_tabletopic_configs(commit_interval_ms) + self._run_perf(topic_prefix, topics=1, partitions=16, send_rate=send_rate, + commit_interval_ms=commit_interval_ms, duration_seconds=test_duration_seconds, + value_schema_json=value_schema_json, topic_configs=topic_configs) + created_topic = self._perf_topic_name(topic_prefix, partitions=16, index=0) + produced_offsets = int(self._sum_kafka_end_offsets_stable(created_topic, settle_window_sec=2, timeout_sec=15)) + self.logger.info(f"Produced messages by offsets(stable): {produced_offsets} (target ~{target_msgs})") + return created_topic, produced_offsets + + def _perf_topic_name(self, prefix: str, partitions: int = 16, index: int = 0) -> str: + return f"__automq_perf_{prefix}_{partitions:04d}_{index:07d}" + + @parametrize(version=str(DEV_BRANCH), interbroker_security_protocol=SecurityConfig.PLAINTEXT) + @parametrize(version=str(DEV_BRANCH), interbroker_security_protocol=SecurityConfig.SSL) + def test_tabletopic_avro_e2e_flow(self, version, interbroker_security_protocol): + """ + Tests the end-to-end flow of the table topic feature with Avro messages. + """ + self.configure_security( + security_protocol=SecurityConfig.PLAINTEXT, + interbroker_security_protocol=interbroker_security_protocol, + restart_cluster=True) + topic_name = f"{self.base_topic}_avro" + target_msgs = 5000 + + commit_interval_ms = 2000 + # Step 1: Produce Avro messages via AutoMQPerformanceService (random Avro if no values-file) + created_topic, produced = self._produce_avro_data_with_perf(topic_name, json.dumps(self.avro_schema), + target_msgs=target_msgs, commit_interval_ms=commit_interval_ms) + # Step 2: Verify table exists and contains expected data (polling for commit interval) + self._verify_data_in_iceberg(created_topic, produced, commit_interval_ms=commit_interval_ms) diff --git a/tests/kafkatest/tests/core/automq_tabletopic_matrix_test.py b/tests/kafkatest/tests/core/automq_tabletopic_matrix_test.py new file mode 100644 index 0000000000..ae4c84e4aa --- /dev/null +++ b/tests/kafkatest/tests/core/automq_tabletopic_matrix_test.py @@ -0,0 +1,75 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from ducktape.mark import parametrize +from kafkatest.version import DEV_BRANCH +from .automq_tabletopic_base import TableTopicBase + + +class TableTopicMatrixTest(TableTopicBase): + def __init__(self, test_context): + super(TableTopicMatrixTest, self).__init__(test_context) + + # Simple Avro schema + self.avro_schema = { + "type": "record", + "name": "MatrixAvroSchema", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "ts", "type": "long"} + ] + } + # Use TableTopicBase helpers for perf, offsets, and Iceberg verification + + # Use TableTopicBase._run_perf + + @parametrize(commit_interval_ms=500) + @parametrize(commit_interval_ms=2000) + @parametrize(commit_interval_ms=10000) + def test_commit_interval_matrix(self, commit_interval_ms, version=str(DEV_BRANCH)): + topic_prefix = f"tt_ci_{commit_interval_ms}" + topics = 1 + partitions = 16 + send_rate = 200 + duration_seconds = 60 + topic_configs = self.default_tabletopic_configs(commit_interval_ms) + self._run_perf(topic_prefix, topics, partitions, send_rate, commit_interval_ms, duration_seconds, + value_schema_json=json.dumps(self.avro_schema), + topic_configs=topic_configs) + # Single topic + topic_name = self._perf_topic_name(topic_prefix, partitions, 0) + produced = int(self._sum_kafka_end_offsets_stable(topic_name, 2, 15)) + self._verify_data_in_iceberg(topic_name, produced, commit_interval_ms=commit_interval_ms) + + @parametrize(topic_count=1) + @parametrize(topic_count=10) + @parametrize(topic_count=100) + def test_topics_count_matrix(self, topic_count, version=str(DEV_BRANCH)): + topic_prefix = f"tt_topics_{topic_count}" + # Total 1000 partitions across all topics, total write rate 10000 across all partitions + partitions = 1000 // topic_count # Distribute total partitions across topics + send_rate = (10000 // 1000) * partitions # Rate per partition (10) * partitions per topic + commit_interval_ms = 2000 + topic_configs = self.default_tabletopic_configs(commit_interval_ms) + duration_seconds = 60 + self._run_perf(topic_prefix, topic_count, partitions, send_rate, commit_interval_ms, duration_seconds, + value_schema_json=json.dumps(self.avro_schema), + topic_configs=topic_configs) + for i in range(topic_count): + name = self._perf_topic_name(topic_prefix, partitions, i) + produced = int(self._sum_kafka_end_offsets_stable(name, 2, 30)) + self._verify_data_in_iceberg(name, produced, commit_interval_ms=commit_interval_ms) diff --git a/tests/kafkatest/tests/core/automq_tabletopic_partition_reassignment_test.py b/tests/kafkatest/tests/core/automq_tabletopic_partition_reassignment_test.py new file mode 100644 index 0000000000..88b943f33a --- /dev/null +++ b/tests/kafkatest/tests/core/automq_tabletopic_partition_reassignment_test.py @@ -0,0 +1,93 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import random +import time + +from ducktape.utils.util import wait_until +from .automq_tabletopic_base import TableTopicBase + + +class TableTopicPartitionReassignmentTest(TableTopicBase): + def __init__(self, test_context): + super(TableTopicPartitionReassignmentTest, self).__init__(test_context) + + self.avro_schema = { + "type": "record", + "name": "ReassignAvroSchema", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "ts", "type": "long"} + ] + } + + def _wait_topic_exists(self, topic_name: str, timeout_sec: int = 60): + def exists(): + for t in self.kafka.list_topics(): + if t == topic_name: + return True + return False + wait_until(exists, timeout_sec=timeout_sec, err_msg=f"Topic {topic_name} not found in {timeout_sec}s") + + def _reassign_partitions(self, topic: str, num_partitions: int): + desc = self.kafka.describe_topic(topic) + info = self.kafka.parse_describe_topic(desc) + self.logger.info(f"Before reassignment: {info}") + # Shuffle partition mapping: swap partition ids while keeping replicas lists + shuffled = list(range(0, num_partitions)) + random.shuffle(shuffled) + for i in range(0, num_partitions): + info["partitions"][i]["partition"] = shuffled[i] + self.logger.info(f"Reassignment plan: {info}") + # Execute and wait for completion + self.kafka.execute_reassign_partitions(info) + wait_until(lambda: self.kafka.verify_reassign_partitions(info), timeout_sec=120, backoff_sec=1, + err_msg="Partition reassignment did not complete in time") + + def test_partition_reassignment_during_produce(self): + topic_prefix = "tt_reassign" + partitions = 16 + send_rate = 200 + commit_interval_ms = 2000 + test_duration_seconds = 120 + + topic_configs = self.default_tabletopic_configs(commit_interval_ms) + + try: + # Kick off a delayed reassignment via a side thread + def do_reassign(): + topic_name = self._perf_topic_name(topic_prefix, partitions, 0) + self._wait_topic_exists(topic_name, 60) + time.sleep(5) # ensure leaders ready + self._reassign_partitions(topic_name, partitions) + + import threading + th = threading.Thread(target=lambda: (time.sleep(30), do_reassign()), daemon=True) + th.start() + + # Run perf using base helper (blocks for duration) + self._run_perf(topic_prefix, 1, partitions, send_rate, + commit_interval_ms, test_duration_seconds, + value_schema_json=json.dumps(self.avro_schema), + topic_configs=topic_configs) + th.join(timeout=5) + finally: + pass + + topic_name = self._perf_topic_name(topic_prefix, partitions, 0) + produced = int(self._sum_kafka_end_offsets_stable(topic_name, 3, 60)) + self._verify_data_in_iceberg(topic_name, produced, commit_interval_ms=commit_interval_ms) diff --git a/tests/kafkatest/tests/core/automq_tabletopic_schema_evolution_test.py b/tests/kafkatest/tests/core/automq_tabletopic_schema_evolution_test.py new file mode 100644 index 0000000000..99c6a4db13 --- /dev/null +++ b/tests/kafkatest/tests/core/automq_tabletopic_schema_evolution_test.py @@ -0,0 +1,103 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import time +from typing import Dict + +from .automq_tabletopic_base import TableTopicBase + +class TableTopicSchemaEvolutionTest(TableTopicBase): + def __init__(self, test_context): + super(TableTopicSchemaEvolutionTest, self).__init__(test_context) + + # Define 4 compatible schemas: each adds a field with a default value + self.schemas = [ + { + "type": "record", + "name": "SchemaV1", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"} + ] + }, + { + "type": "record", + "name": "SchemaV2", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "email", "type": "string", "default": ""} + ] + }, + { + "type": "record", + "name": "SchemaV3", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "email", "type": "string", "default": ""}, + {"name": "age", "type": "int", "default": 0} + ] + }, + { + "type": "record", + "name": "SchemaV4", + "fields": [ + {"name": "id", "type": "long"}, + {"name": "name", "type": "string"}, + {"name": "email", "type": "string", "default": ""}, + {"name": "age", "type": "int", "default": 0}, + {"name": "country", "type": "string", "default": "cn"} + ] + } + ] + + def _run_phase(self, topic_prefix: str, schema_json: str, partitions: int, send_rate: int, + commit_interval_ms: int, duration_seconds: int = 60, topic_configs: Dict[str, str] = None): + # Allow caller to provide topic_configs; default to base helper if not provided + effective_topic_configs = topic_configs if topic_configs is not None else self.default_tabletopic_configs(commit_interval_ms) + self._run_perf(topic_prefix, 1, partitions, send_rate, + commit_interval_ms, duration_seconds, + value_schema_json=schema_json, + topic_configs=effective_topic_configs) + + def test_schema_evolution_4_phases(self): + topic_prefix = "tt_schema_evo" + partitions = 16 + send_rate = 200 + commit_interval_ms = 2000 + topic_configs = self.default_tabletopic_configs(commit_interval_ms) + + # Phase A: v1 + self._run_phase(topic_prefix, json.dumps(self.schemas[0]), partitions, send_rate, commit_interval_ms, 60, topic_configs=topic_configs) + topic_name = self._perf_topic_name(topic_prefix, partitions, 0) + produced_a = int(self._sum_kafka_end_offsets_stable(topic_name, 2, 30)) + self._verify_data_in_iceberg(topic_name, produced_a, commit_interval_ms) + + # Phase B: v2 + self._run_phase(topic_prefix, json.dumps(self.schemas[1]), partitions, send_rate, commit_interval_ms, 60, topic_configs=topic_configs) + produced_b = int(self._sum_kafka_end_offsets_stable(topic_name, 2, 30)) + self._verify_data_in_iceberg(topic_name, produced_b, commit_interval_ms) + + # Phase C: v3 + self._run_phase(topic_prefix, json.dumps(self.schemas[2]), partitions, send_rate, commit_interval_ms, 60, topic_configs=topic_configs) + produced_c = int(self._sum_kafka_end_offsets_stable(topic_name, 2, 30)) + self._verify_data_in_iceberg(topic_name, produced_c, commit_interval_ms) + + # Phase D: v4 + self._run_phase(topic_prefix, json.dumps(self.schemas[3]), partitions, send_rate, commit_interval_ms, 60, topic_configs=topic_configs) + produced_d = int(self._sum_kafka_end_offsets_stable(topic_name, 2, 30)) + self._verify_data_in_iceberg(topic_name, produced_d, commit_interval_ms) diff --git a/tests/kafkatest/tests/core/automq_telemetry_test.py b/tests/kafkatest/tests/core/automq_telemetry_test.py new file mode 100644 index 0000000000..1b7ea6942d --- /dev/null +++ b/tests/kafkatest/tests/core/automq_telemetry_test.py @@ -0,0 +1,280 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import subprocess +import time + +from ducktape.mark.resource import cluster +from ducktape.mark import parametrize +from ducktape.tests.test import Test +from ducktape.utils.util import wait_until + +from kafkatest.services.kafka import KafkaService, quorum +from kafkatest.services.security.security_config import SecurityConfig +from kafkatest.services.verifiable_producer import VerifiableProducer +from kafkatest.services.zookeeper import ZookeeperService + + +class AutoMQBrokerTelemetryTest(Test): + """End-to-end validation for AutoMQ telemetry and log uploader integration in the broker.""" + + TOPIC = "automq-telemetry-topic" + + def __init__(self, test_context): + super(AutoMQBrokerTelemetryTest, self).__init__(test_context) + self.num_brokers = 1 + self.zk = None + self.kafka = None + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _start_kafka(self, server_overrides=None, per_node_overrides=None, extra_env=None): + if quorum.for_test(self.test_context) == quorum.zk and self.zk is None: + self.zk = ZookeeperService(self.test_context, 1) + self.zk.start() + + self.kafka = KafkaService( + self.test_context, + self.num_brokers, + self.zk, + security_protocol=SecurityConfig.PLAINTEXT, + topics={}, + server_prop_overrides=server_overrides, + per_node_server_prop_overrides=per_node_overrides, + extra_env=extra_env, + ) + + self.kafka.start() + self.kafka.create_topic({ + "topic": self.TOPIC, + "partitions": 1, + "replication-factor": 1, + }) + + def _stop_kafka(self): + if self.kafka is not None: + self.kafka.stop() + self.kafka = None + if self.zk is not None: + self.zk.stop() + self.zk = None + + def _produce_messages(self, max_messages=200, throughput=1000): + producer = VerifiableProducer( + self.test_context, + num_nodes=1, + kafka=self.kafka, + topic=self.TOPIC, + max_messages=max_messages, + throughput=throughput, + ) + producer.start() + try: + wait_until( + lambda: producer.num_acked >= max_messages, + timeout_sec=60, + backoff_sec=5, + err_msg="Producer failed to deliver expected number of messages", + ) + finally: + try: + producer.stop() + except Exception as e: + self.logger.warn("Error stopping producer: %s", e) + + def _metrics_ready(self, node, port): + try: + cmd = f"curl -sf http://localhost:{port}/metrics" + output = "".join(list(node.account.ssh_capture(cmd, allow_fail=True))) + return bool(output.strip()) + except Exception: + return False + + def _wait_for_metrics_available(self, port=9464, timeout_sec=90): + for node in self.kafka.nodes: + wait_until( + lambda n=node: self._metrics_ready(n, port), + timeout_sec=timeout_sec, + backoff_sec=5, + err_msg=f"Metrics endpoint not available on {node.account.hostname}", + ) + + def _fetch_metrics(self, node, port=9464): + cmd = f"curl -sf http://localhost:{port}/metrics" + return "".join(list(node.account.ssh_capture(cmd, allow_fail=True))) + + def _assert_prometheus_metrics(self, metrics_output, expected_labels=None): + assert metrics_output.strip(), "Metrics endpoint returned no data" + + metric_lines = [ + line for line in metrics_output.splitlines() + if line.strip() and not line.startswith('#') + ] + assert metric_lines, "No metric datapoints found in Prometheus output" + + kafka_lines = [line for line in metric_lines if 'kafka_' in line or 'automq' in line] + assert kafka_lines, "Expected broker metrics not present in Prometheus output" + + if expected_labels: + for label in expected_labels: + assert label in metrics_output, f"Expected label '{label}' absent from metrics output" + + if "# HELP" not in metrics_output and "# TYPE" not in metrics_output: + self.logger.warning("Metrics output missing HELP/TYPE comments – format may not follow Prometheus conventions") + + def _list_s3_objects(self, prefix): + objects, _ = self.kafka.get_bucket_objects() + return [obj for obj in objects if obj["path"].startswith(prefix)] + + def _clear_s3_prefix(self, bucket, prefix): + cmd = f"aws s3 rm s3://{bucket}/{prefix} --recursive --endpoint=http://10.5.0.2:4566" + ret, out = subprocess.getstatusoutput(cmd) + if ret != 0: + self.logger.info("Ignoring cleanup error for prefix %s: %s", prefix, out) + + def _check_port_listening(self, node, port): + """Check if a port is listening on the given node.""" + try: + result = list(node.account.ssh_capture(f"netstat -ln | grep :{port}", allow_fail=True)) + return len(result) > 0 + except Exception: + return False + + def _extract_metric_samples(self, metrics_output, metric_name): + samples = [] + for line in metrics_output.splitlines(): + if line.startswith(metric_name): + parts = line.split() + if len(parts) >= 2: + try: + samples.append(float(parts[-1])) + except ValueError: + continue + return samples + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + @cluster(num_nodes=4) + def test_prometheus_metrics_exporter(self): + """Verify that the broker exposes Prometheus metrics via the AutoMQ OpenTelemetry module.""" + cluster_label = f"kafka-core-prom-{int(time.time())}" + server_overrides = [ + ["s3.telemetry.metrics.exporter.uri", "prometheus://0.0.0.0:9464"], + ["s3.telemetry.exporter.report.interval.ms", "10000"], + ["s3.telemetry.metrics.base.labels", "component=broker"] + ] + + self._start_kafka(server_overrides=server_overrides) + + try: + self._produce_messages(max_messages=200) + self._wait_for_metrics_available() + + for node in self.kafka.nodes: + output = self._fetch_metrics(node) + self._assert_prometheus_metrics( + output, + expected_labels=['instance="'] + ) + finally: + self._stop_kafka() + + @cluster(num_nodes=4) + def test_s3_metrics_exporter(self): + """Verify that broker metrics are exported to S3 via the AutoMQ telemetry module.""" + bucket_name = "ko3" + metrics_prefix = f"automq/metrics" + + self._clear_s3_prefix(bucket_name, metrics_prefix) + + server_overrides = [ + ["s3.telemetry.metrics.exporter.uri", f"ops://{bucket_name}"], + ["s3.telemetry.ops.enabled", "true"], + ["s3.ops.buckets", f"0@s3://{bucket_name}?endpoint=http://10.5.0.2:4566®ion=us-east-1"], + ] + + original_num_brokers = self.num_brokers + node_count = max(2, self.num_brokers) + self.num_brokers = node_count + + try: + self._start_kafka(server_overrides=server_overrides) + + self._produce_messages(max_messages=200) + + def _metrics_uploaded(): + objects = self._list_s3_objects(metrics_prefix) + if objects: + self.logger.info("Found %d metrics objects for prefix %s", len(objects), metrics_prefix) + return len(objects) > 0 + + wait_until( + _metrics_uploaded, + timeout_sec=180, + backoff_sec=10, + err_msg="Timed out waiting for S3 metrics export" + ) + + finally: + self.num_brokers = original_num_brokers + self._stop_kafka() + + @cluster(num_nodes=4) + def test_s3_log_uploader(self): + """Verify that broker logs are uploaded to S3 via the AutoMQ log uploader module.""" + bucket_name = "ko3" + logs_prefix = f"automq/logs" + + self._clear_s3_prefix(bucket_name, logs_prefix) + + server_overrides = [ + ["s3.telemetry.ops.enabled", "true"], + ["s3.ops.buckets", f"0@s3://{bucket_name}?endpoint=http://10.5.0.2:4566®ion=us-east-1"], + ] + + extra_env = [ + "AUTOMQ_OBSERVABILITY_UPLOAD_INTERVAL=15000", + "AUTOMQ_OBSERVABILITY_CLEANUP_INTERVAL=60000" + ] + + original_num_brokers = self.num_brokers + node_count = max(2, self.num_brokers) + self.num_brokers = node_count + + try: + self._start_kafka(server_overrides=server_overrides, extra_env=extra_env) + + self._produce_messages(max_messages=300) + + def _logs_uploaded(): + objects = self._list_s3_objects(logs_prefix) + if objects: + self.logger.info("Found %d log objects for prefix %s", len(objects), logs_prefix) + return len(objects) > 0 + + wait_until( + _logs_uploaded, + timeout_sec=240, + backoff_sec=15, + err_msg="Timed out waiting for S3 log upload" + ) + + finally: + self.num_brokers = original_num_brokers + self._stop_kafka() diff --git a/tests/kafkatest/tests/core/group_mode_transactions_test.py b/tests/kafkatest/tests/core/group_mode_transactions_test.py index cf9c9c5bf8..52f5099926 100644 --- a/tests/kafkatest/tests/core/group_mode_transactions_test.py +++ b/tests/kafkatest/tests/core/group_mode_transactions_test.py @@ -49,7 +49,7 @@ def __init__(self, test_context): self.num_input_partitions = 9 self.num_output_partitions = 9 self.num_copiers = 3 - self.num_seed_messages = 100000 + self.num_seed_messages = 1000 self.transaction_size = 750 # The transaction timeout should be lower than the progress timeout, but at # least as high as the request timeout (which is 30s by default). When the diff --git a/tests/kafkatest/tests/core/reassign_partitions_test.py b/tests/kafkatest/tests/core/reassign_partitions_test.py index 93bd3f436e..a3b2b0c55e 100644 --- a/tests/kafkatest/tests/core/reassign_partitions_test.py +++ b/tests/kafkatest/tests/core/reassign_partitions_test.py @@ -61,7 +61,7 @@ def __init__(self, test_context): }, controller_num_nodes_override=self.num_zk) self.timeout_sec = 60 - self.producer_throughput = 1000 + self.producer_throughput = 10000 self.num_producers = 1 self.num_consumers = 1 @@ -94,6 +94,10 @@ def reassign_partitions(self, bounce_brokers): partition_info["partitions"][i]["partition"] = shuffled_list[i] self.logger.debug("Jumbled partitions: " + str(partition_info)) + acked_partitions = self.producer.acked_by_partition + for partition in acked_partitions: + self.logger.debug("Partition acked " + str(partition)) + def check_all_partitions(): acked_partitions = self.producer.acked_by_partition for i in range(self.num_partitions): @@ -103,8 +107,8 @@ def check_all_partitions(): # ensure all partitions have data so we don't hit OutOfOrderExceptions due to broker restarts wait_until(check_all_partitions, - timeout_sec=60, - err_msg="Failed to produce to all partitions in 30s") + timeout_sec=120, + err_msg="Failed to produce to all partitions in 120s") # send reassign partitions command self.kafka.execute_reassign_partitions(partition_info) diff --git a/tests/kafkatest/tests/core/replica_scale_test.py b/tests/kafkatest/tests/core/replica_scale_test.py index 8bc8648731..b9d91af37e 100644 --- a/tests/kafkatest/tests/core/replica_scale_test.py +++ b/tests/kafkatest/tests/core/replica_scale_test.py @@ -98,10 +98,10 @@ def test_produce_consume(self, topic_count, partition_count, replication_factor, produce_spec = ProduceBenchWorkloadSpec(0, TaskSpec.MAX_DURATION_MS, producer_workload_service.producer_node, producer_workload_service.bootstrap_servers, - target_messages_per_sec=150000, + target_messages_per_sec=1500, # optimize multiple partition read # max_messages=3400000, - max_messages=1700000, + max_messages=17000, producer_conf={}, admin_client_conf={}, common_client_conf={}, @@ -117,8 +117,8 @@ def test_produce_consume(self, topic_count, partition_count, replication_factor, consume_spec = ConsumeBenchWorkloadSpec(0, TaskSpec.MAX_DURATION_MS, consumer_workload_service.consumer_node, consumer_workload_service.bootstrap_servers, - target_messages_per_sec=150000, - max_messages=1700000, + target_messages_per_sec=1500, + max_messages=17000, consumer_conf=consumer_conf, admin_client_conf={}, common_client_conf={}, diff --git a/tests/kafkatest/tests/core/round_trip_fault_test.py b/tests/kafkatest/tests/core/round_trip_fault_test.py index a22100f24f..80637734cc 100644 --- a/tests/kafkatest/tests/core/round_trip_fault_test.py +++ b/tests/kafkatest/tests/core/round_trip_fault_test.py @@ -51,8 +51,8 @@ def __init__(self, test_context): self.round_trip_spec = RoundTripWorkloadSpec(0, TaskSpec.MAX_DURATION_MS, self.workload_service.client_node, self.workload_service.bootstrap_servers, - target_messages_per_sec=10000, - max_messages=100000, + target_messages_per_sec=1000, + max_messages=10000, active_topics=active_topics) def setUp(self): diff --git a/tests/kafkatest/tests/core/transactions_test.py b/tests/kafkatest/tests/core/transactions_test.py index 12191ff68c..7433c2fd94 100644 --- a/tests/kafkatest/tests/core/transactions_test.py +++ b/tests/kafkatest/tests/core/transactions_test.py @@ -46,7 +46,7 @@ def __init__(self, test_context): # Test parameters self.num_input_partitions = 2 self.num_output_partitions = 3 - self.num_seed_messages = 100000 + self.num_seed_messages = 10000 self.transaction_size = 750 # The transaction timeout should be lower than the progress timeout, but at @@ -189,7 +189,7 @@ def copy_messages_transactionally(self, failure_mode, bounce_target, elif bounce_target == "clients": self.bounce_copiers(copiers, clean_shutdown) - copier_timeout_sec = 120 + copier_timeout_sec = 600 for copier in copiers: wait_until(lambda: copier.is_done, timeout_sec=copier_timeout_sec, @@ -218,22 +218,6 @@ def setup_topics(self): } @cluster(num_nodes=9) - # @matrix( - # failure_mode=["hard_bounce", "clean_bounce"], - # bounce_target=["brokers", "clients"], - # check_order=[True, False], - # use_group_metadata=[True, False], - # metadata_quorum=[quorum.zk], - # use_new_coordinator=[False] - # ) - @matrix( - failure_mode=["hard_bounce", "clean_bounce"], - bounce_target=["brokers", "clients"], - check_order=[True, False], - use_group_metadata=[True, False], - metadata_quorum=quorum.all_kraft, - use_new_coordinator=[False] - ) @matrix( failure_mode=["hard_bounce", "clean_bounce"], bounce_target=["brokers", "clients"], diff --git a/tests/kafkatest/tests/produce_consume_validate.py b/tests/kafkatest/tests/produce_consume_validate.py index c691cbc12b..2c01d0db99 100644 --- a/tests/kafkatest/tests/produce_consume_validate.py +++ b/tests/kafkatest/tests/produce_consume_validate.py @@ -31,10 +31,10 @@ def __init__(self, test_context): super(ProduceConsumeValidateTest, self).__init__(test_context=test_context) # How long to wait for the producer to declare itself healthy? This can # be overidden by inheriting classes. - self.producer_start_timeout_sec = 20 + self.producer_start_timeout_sec = 600 # How long to wait for the consumer to start consuming messages? - self.consumer_start_timeout_sec = 60 + self.consumer_start_timeout_sec = 600 # How long wait for the consumer process to fork? This # is important in the case when the consumer is starting from the end, @@ -65,7 +65,7 @@ def start_producer_and_consumer(self): self.producer.start() - wait_until(lambda: self.producer.num_acked > 5, + wait_until(lambda: self.producer.num_acked > 19, timeout_sec=self.producer_start_timeout_sec, err_msg="Producer failed to produce messages for %ds." %\ self.producer_start_timeout_sec) diff --git a/tests/kafkatest/tests/streams/base_streams_test.py b/tests/kafkatest/tests/streams/base_streams_test.py index a7c7a7b10d..eb894d1b6d 100644 --- a/tests/kafkatest/tests/streams/base_streams_test.py +++ b/tests/kafkatest/tests/streams/base_streams_test.py @@ -92,7 +92,7 @@ def get_configs(extra_configs=""): def wait_for_verification(self, processor, message, file, num_lines=1): wait_until(lambda: self.verify_from_file(processor, message, file) >= num_lines, - timeout_sec=60, + timeout_sec=300, err_msg="Did expect to read '%s' from %s" % (message, processor.node.account)) def verify_from_file(self, processor, message, file): diff --git a/tests/kafkatest/tests/streams/streams_broker_down_resilience_test.py b/tests/kafkatest/tests/streams/streams_broker_down_resilience_test.py index 9eeeea3c98..3f8bac862b 100644 --- a/tests/kafkatest/tests/streams/streams_broker_down_resilience_test.py +++ b/tests/kafkatest/tests/streams/streams_broker_down_resilience_test.py @@ -253,7 +253,7 @@ def test_streams_should_failover_while_brokers_down(self, metadata_quorum, use_n processor_3.start() monitor.wait_until(rebalance, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw output '%s' on " % rebalance) + str(processor_3.node.account)) with processor.node.account.monitor_log(processor.STDOUT_FILE) as monitor_1: @@ -263,23 +263,23 @@ def test_streams_should_failover_while_brokers_down(self, metadata_quorum, use_n self.assert_produce(self.inputTopic, "sending_message_after_normal_broker_start", num_messages=self.num_messages, - timeout_sec=120) + timeout_sec=300) monitor_1.wait_until(self.message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.message) + str(processor.node.account)) monitor_2.wait_until(self.message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.message) + str(processor_2.node.account)) monitor_3.wait_until(self.message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.message) + str(processor_3.node.account)) self.assert_consume(self.client_id, "consuming_message_after_normal_broker_start", self.outputTopic, num_messages=self.num_messages, - timeout_sec=120) + timeout_sec=300) node = self.kafka.leader(self.inputTopic) self.kafka.stop_node(node) @@ -294,13 +294,13 @@ def test_streams_should_failover_while_brokers_down(self, metadata_quorum, use_n self.kafka.start_node(node) monitor_1.wait_until(self.connected_message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.connected_message) + str(processor.node.account)) monitor_2.wait_until(self.connected_message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.connected_message) + str(processor_2.node.account)) monitor_3.wait_until(self.connected_message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.connected_message) + str(processor_3.node.account)) with processor.node.account.monitor_log(processor.STDOUT_FILE) as monitor_1: @@ -310,21 +310,21 @@ def test_streams_should_failover_while_brokers_down(self, metadata_quorum, use_n self.assert_produce(self.inputTopic, "sending_message_after_hard_bouncing_streams_instance_bouncing_broker", num_messages=self.num_messages, - timeout_sec=120) + timeout_sec=300) monitor_1.wait_until(self.message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.message) + str(processor.node.account)) monitor_2.wait_until(self.message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.message) + str(processor_2.node.account)) monitor_3.wait_until(self.message, - timeout_sec=120, + timeout_sec=300, err_msg=("Never saw '%s' on " % self.message) + str(processor_3.node.account)) self.assert_consume(self.client_id, "consuming_message_after_stopping_streams_instance_bouncing_broker", self.outputTopic, num_messages=self.num_messages, - timeout_sec=120) + timeout_sec=300) self.kafka.stop() diff --git a/tests/kafkatest/tests/streams/streams_relational_smoke_test.py b/tests/kafkatest/tests/streams/streams_relational_smoke_test.py index c53715d769..db529ce899 100644 --- a/tests/kafkatest/tests/streams/streams_relational_smoke_test.py +++ b/tests/kafkatest/tests/streams/streams_relational_smoke_test.py @@ -64,8 +64,8 @@ def start_node(self, node): raise RuntimeError("No process ids recorded") def await_command(self, command): - wait_until(lambda: self.node.account.ssh(command, allow_fail=True), - timeout_sec=60, + wait_until(lambda: self.node.account.ssh(command, allow_fail=True) == 0, + timeout_sec=300, err_msg="Command [%s] never passed in the timeout" ) diff --git a/tests/kafkatest/tests/streams/streams_standby_replica_test.py b/tests/kafkatest/tests/streams/streams_standby_replica_test.py index d9804408df..960cfa776d 100644 --- a/tests/kafkatest/tests/streams/streams_standby_replica_test.py +++ b/tests/kafkatest/tests/streams/streams_standby_replica_test.py @@ -120,12 +120,12 @@ def test_standby_tasks_rebalance(self, metadata_quorum, use_new_coordinator=Fals self.wait_for_verification(processor_3, "ACTIVE_TASKS:2 STANDBY_TASKS:[1-3]", processor_3.STDOUT_FILE) self.assert_consume(self.client_id, "assert all messages consumed from %s" % self.streams_sink_topic_1, - self.streams_sink_topic_1, self.num_messages) + self.streams_sink_topic_1, self.num_messages, timeout_sec=300) self.assert_consume(self.client_id, "assert all messages consumed from %s" % self.streams_sink_topic_2, - self.streams_sink_topic_2, self.num_messages) + self.streams_sink_topic_2, self.num_messages, timeout_sec=300) wait_until(lambda: producer.num_acked >= self.num_messages, - timeout_sec=60, + timeout_sec=300, err_msg="Failed to send all %s messages" % str(self.num_messages)) producer.stop() diff --git a/tests/suites/connect_enterprise_test_suite1.yml b/tests/suites/connect_enterprise_test_suite1.yml new file mode 100644 index 0000000000..54c440b6a0 --- /dev/null +++ b/tests/suites/connect_enterprise_test_suite1.yml @@ -0,0 +1,6 @@ +connect_enterprise_test_suite: + included: + - ../kafkatest/tests/connect/connect_remote_write_test.py::ConnectRemoteWriteTest.test_opentelemetry_remote_write_exporter + - ../kafkatest/tests/connect/connect_remote_write_test.py::ConnectRemoteWriteTest.test_remote_write_with_compression + - ../kafkatest/tests/connect/connect_remote_write_test.py::ConnectRemoteWriteTest.test_remote_write_batch_size_limits + - ../kafkatest/tests/connect/connect_remote_write_test.py::ConnectRemoteWriteTest.test_remote_write_server_unavailable diff --git a/tests/suites/connect_test_suite2.yml b/tests/suites/connect_test_suite2.yml index 7a55799851..4b267fdcb5 100644 --- a/tests/suites/connect_test_suite2.yml +++ b/tests/suites/connect_test_suite2.yml @@ -3,4 +3,8 @@ connect_test_suite: - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_restart_failed_connector - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_restart_failed_task - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_restart_connector_and_tasks_failed_connector - - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_restart_connector_and_tasks_failed_task \ No newline at end of file + - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_restart_connector_and_tasks_failed_task + - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_opentelemetry_metrics_basic + - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_opentelemetry_metrics_comprehensive + - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_metrics_under_load + - ../kafkatest/tests/connect/connect_distributed_test.py::ConnectDistributedTest.test_opentelemetry_s3_metrics_exporter diff --git a/tests/suites/main_enterprise_test_suite1.yml b/tests/suites/main_enterprise_test_suite1.yml new file mode 100644 index 0000000000..19c1e7d19c --- /dev/null +++ b/tests/suites/main_enterprise_test_suite1.yml @@ -0,0 +1,7 @@ +main_enterprise_test_suite: + included: + - ../kafkatest/tests/core/automq_remote_write_test.py::AutoMQRemoteWriteTest.test_remote_write_metrics_exporter + - ../kafkatest/tests/core/automq_remote_write_test.py::AutoMQRemoteWriteTest.test_remote_write_with_compression + - ../kafkatest/tests/core/automq_remote_write_test.py::AutoMQRemoteWriteTest.test_remote_write_batch_size_limits + - ../kafkatest/tests/core/automq_remote_write_test.py::AutoMQRemoteWriteTest.test_remote_write_server_unavailable + diff --git a/tests/suites/main_kos_test_suite2.yml b/tests/suites/main_kos_test_suite2.yml index d3cc9dd37a..4e72a57e72 100644 --- a/tests/suites/main_kos_test_suite2.yml +++ b/tests/suites/main_kos_test_suite2.yml @@ -33,6 +33,7 @@ core_test_suite: - ../kafkatest/tests/core/mirror_maker_test.py - ../kafkatest/tests/core/produce_bench_test.py - ../kafkatest/tests/core/transactions_test.py + - ../kafkatest/tests/core/automq_remote_write_test.py # cannot downgrade to official kafka - ../kafkatest/tests/core/downgrade_test.py::TestDowngrade.test_upgrade_and_downgrade diff --git a/tests/suites/main_kos_test_suite4.yml b/tests/suites/main_kos_test_suite4.yml index e7da9746e8..8f9e7ec5e5 100644 --- a/tests/suites/main_kos_test_suite4.yml +++ b/tests/suites/main_kos_test_suite4.yml @@ -17,3 +17,7 @@ core_test_suite: included: - ../kafkatest/tests/core/transactions_test.py + - ../kafkatest/tests/core/automq_telemetry_test.py::AutoMQBrokerTelemetryTest.test_prometheus_metrics_exporter + - ../kafkatest/tests/core/automq_telemetry_test.py::AutoMQBrokerTelemetryTest.test_prometheus_metrics_under_load + - ../kafkatest/tests/core/automq_telemetry_test.py::AutoMQBrokerTelemetryTest.test_s3_metrics_exporter + - ../kafkatest/tests/core/automq_telemetry_test.py::AutoMQBrokerTelemetryTest.test_s3_log_uploader diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/PerfCommand.java b/tools/src/main/java/org/apache/kafka/tools/automq/PerfCommand.java index 1145489522..0f26c6cb54 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/PerfCommand.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/PerfCommand.java @@ -1,16 +1,25 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.tools.automq.perf.ConsumerService; import org.apache.kafka.tools.automq.perf.PerfConfig; import org.apache.kafka.tools.automq.perf.ProducerService; @@ -23,19 +32,38 @@ import com.automq.stream.s3.metrics.TimerUtil; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.ObjectWriter; +import com.google.common.base.Strings; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; import java.text.SimpleDateFormat; import java.time.Duration; import java.util.ArrayList; +import java.util.Collections; import java.util.Date; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; +import java.util.function.Function; + +import io.confluent.kafka.schemaregistry.avro.AvroSchema; +import io.confluent.kafka.schemaregistry.avro.AvroSchemaUtils; +import io.confluent.kafka.serializers.KafkaAvroSerializer; import static org.apache.kafka.tools.automq.perf.StatsCollector.printAndCollectStats; @@ -51,7 +79,13 @@ public class PerfCommand implements AutoCloseable { private final ProducerService producerService; private final ConsumerService consumerService; private final Stats stats = new Stats(); + /** + * Partitions that are ready to be consumed. + * Only used during the initial topic readiness check, which is, {@link #preparing} is true. + */ + private final Set readyPartitions = Collections.newSetFromMap(new ConcurrentHashMap<>()); + private volatile boolean preparing = true; private volatile boolean running = true; public static void main(String[] args) throws Exception { @@ -65,7 +99,7 @@ private PerfCommand(PerfConfig config) { this.config = config; this.topicService = new TopicService(config.bootstrapServer(), config.adminConfig()); this.producerService = new ProducerService(); - this.consumerService = new ConsumerService(config.bootstrapServer()); + this.consumerService = new ConsumerService(config.bootstrapServer(), config.adminConfig()); } private void run() { @@ -73,38 +107,65 @@ private void run() { TimerUtil timer = new TimerUtil(); if (config.reset) { - LOGGER.info("Deleting all topics..."); + LOGGER.info("Deleting all test topics..."); int deleted = topicService.deleteTopics(); - LOGGER.info("Deleted all topics ({} in total), took {} ms", deleted, timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); + LOGGER.info("Deleted all test topics ({} in total), took {} ms", deleted, timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); } - LOGGER.info("Creating topics..."); - List topics = topicService.createTopics(config.topicsConfig()); - LOGGER.info("Created {} topics, took {} ms", topics.size(), timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); + List topics; + if (config.catchupTopicPrefix != null && !config.catchupTopicPrefix.isEmpty()) { + LOGGER.info("Looking for catch-up topics with prefix: {}", config.catchupTopicPrefix); + topics = topicService.findExistingTopicsByPrefix(config.catchupTopicPrefix); + if (topics.isEmpty()) { + throw new RuntimeException("No catch-up topics found with prefix: " + config.catchupTopicPrefix); + } + LOGGER.info("Found {} catch-up topics with prefix '{}'.", topics.size(), config.catchupTopicPrefix); + } else if (config.reuseTopics) { + LOGGER.info("Reusing existing topics with prefix: {}", config.topicPrefix); + topics = topicService.findExistingTopicsByPrefix(config.topicPrefix); + if (topics.isEmpty()) { + LOGGER.warn("No existing topics found with prefix '{}', creating new topics instead.", config.topicPrefix); + topics = topicService.createTopics(config.topicsConfig()); + LOGGER.info("Created {} topics, took {} ms", topics.size(), timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); + } else { + LOGGER.info("Found {} existing topics with prefix '{}'.", topics.size(), config.topicPrefix); + } + } else { + LOGGER.info("Creating topics..."); + topics = topicService.createTopics(config.topicsConfig()); + LOGGER.info("Created {} topics, took {} ms", topics.size(), timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); + } LOGGER.info("Creating consumers..."); - int consumers = consumerService.createConsumers(topics, config.consumersConfig()); - consumerService.start(this::messageReceived); + int consumers = consumerService.createConsumers(topics, config.consumersConfig(), stats); + consumerService.start(this::messageReceived, config.maxConsumeRecordRate); LOGGER.info("Created {} consumers, took {} ms", consumers, timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); LOGGER.info("Creating producers..."); int producers = producerService.createProducers(topics, config.producersConfig(), this::messageSent); LOGGER.info("Created {} producers, took {} ms", producers, timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); - LOGGER.info("Waiting for topics to be ready..."); - waitTopicsReady(); - LOGGER.info("Topics are ready, took {} ms", timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); - - List payloads = randomPayloads(config.recordSize, config.randomRatio, config.randomPoolSize); - producerService.start(payloads, config.sendRate); - - if (config.warmupDurationMinutes > 0) { - LOGGER.info("Warming up for {} minutes...", config.warmupDurationMinutes); - long warmupStart = System.nanoTime(); - long warmupMiddle = warmupStart + TimeUnit.MINUTES.toNanos(config.warmupDurationMinutes) / 2; - producerService.adjustRate(warmupStart, ProducerService.MIN_RATE); - producerService.adjustRate(warmupMiddle, config.sendRate); - collectStats(Duration.ofMinutes(config.warmupDurationMinutes)); + if (config.catchupTopicPrefix != null && !config.catchupTopicPrefix.isEmpty()) { + LOGGER.info("Using catch-up topics, skipping message accumulation phase"); + preparing = false; // Directly start consuming without accumulation + } else { + if (config.awaitTopicReady) { + LOGGER.info("Waiting for topics to be ready..."); + waitTopicsReady(consumerService.consumerCount() > 0); + LOGGER.info("Topics are ready, took {} ms", timer.elapsedAndResetAs(TimeUnit.MILLISECONDS)); + } + Function> payloads = payloads(config, topics); + producerService.start(payloads, config.sendRate); + preparing = false; + + if (config.warmupDurationMinutes > 0) { + LOGGER.info("Warming up for {} minutes...", config.warmupDurationMinutes); + long warmupStart = System.nanoTime(); + long warmupMiddle = warmupStart + TimeUnit.MINUTES.toNanos(config.warmupDurationMinutes) / 2; + producerService.adjustRate(warmupStart, ProducerService.MIN_RATE); + producerService.adjustRate(warmupMiddle, config.sendRate); + collectStats(Duration.ofMinutes(config.warmupDurationMinutes)); + } } Result result; @@ -117,7 +178,21 @@ private void run() { LOGGER.info("Resetting consumer offsets and resuming..."); consumerService.resetOffset(backlogStart, TimeUnit.SECONDS.toMillis(config.groupStartDelaySeconds)); - consumerService.resume(); + + // Select topics for catch-up + int numTopics = topics.size(); + int topicsToResume = (int) Math.ceil(numTopics * (config.consumersDuringCatchupPercentage / 100.0)); + topicsToResume = Math.max(1, Math.min(numTopics, topicsToResume)); + List allTopicNames = topics.stream().map(t -> t.name()).collect(java.util.stream.Collectors.toList()); + java.util.Collections.shuffle(allTopicNames); + List resumeTopics = allTopicNames.subList(0, topicsToResume); + + // No need to pause all again, they're already paused from line 172 + // Just resume the selected topics + consumerService.resumeTopics(resumeTopics); + LOGGER.info("Resuming consumers for topics: {} ({} out of {})", resumeTopics, topicsToResume, numTopics); + // No need to pause specific topics, they're already paused + LOGGER.info("Keeping remaining consumers paused ({} topics)", numTopics - topicsToResume); stats.reset(); producerService.adjustRate(config.sendRateDuringCatchup); @@ -148,23 +223,39 @@ private void messageSent(int size, long sendTimeNanos, Exception exception) { } } - private void messageReceived(byte[] payload, long sendTimeNanos) { - stats.messageReceived(payload.length, sendTimeNanos); + private void messageReceived(TopicPartition topicPartition) { + if (preparing && config.awaitTopicReady && (config.catchupTopicPrefix == null || config.catchupTopicPrefix.isEmpty())) { + readyPartitions.add(topicPartition); + } + } + + private void waitTopicsReady(boolean hasConsumer) { + if (config.catchupTopicPrefix != null && !config.catchupTopicPrefix.isEmpty()) { + LOGGER.info("Using catch-up topics, skipping topic readiness check."); + return; + } + if (hasConsumer) { + waitTopicsReadyWithConsumer(); + } else { + waitTopicsReadyWithoutConsumer(); + } + stats.reset(); } - private void waitTopicsReady() { - int sent = producerService.probe(); + private void waitTopicsReadyWithConsumer() { long start = System.nanoTime(); boolean ready = false; + int expectPartitionCount = config.topics * config.partitionsPerTopic; while (System.nanoTime() < start + TOPIC_READY_TIMEOUT_NANOS) { - long received = stats.toCumulativeStats().totalMessagesReceived; - LOGGER.info("Waiting for topics to be ready... sent: {}, received: {}", sent, received); - if (received >= sent) { + producerService.probe(); + int received = readyPartitions.size(); + LOGGER.info("Waiting for topics to be ready... sent: {}, received: {}", expectPartitionCount, received); + if (received >= expectPartitionCount) { ready = true; break; } try { - Thread.sleep(2000); + Thread.sleep(TimeUnit.SECONDS.toMillis(5)); } catch (InterruptedException e) { throw new RuntimeException(e); } @@ -172,7 +263,36 @@ private void waitTopicsReady() { if (!ready) { throw new RuntimeException("Timeout waiting for topics to be ready"); } - stats.reset(); + } + + private void waitTopicsReadyWithoutConsumer() { + producerService.probe(); + try { + // If there is no consumer, we can only wait for a fixed time to ensure the topic is ready. + Thread.sleep(TimeUnit.SECONDS.toMillis(30)); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + + private Function> payloads(PerfConfig config, List topics) { + if (Strings.isNullOrEmpty(config.valueSchema)) { + List payloads = randomPayloads(config.recordSize, config.randomRatio, config.randomPoolSize); + return topic -> payloads; + } else { + // The producer configs should contain: + // - schema.registry.url: http://localhost:8081 + Map> topic2payloads = new HashMap<>(); + topics.forEach(topic -> { + if (!Strings.isNullOrEmpty(config.valuesFile)) { + topic2payloads.put(topic.name(), schemaPayloads(topic.name(), config.valueSchema, config.valuesFile, config.producerConfigs)); + } else { + // Default: auto-generate random Avro values when no --values-file provided + topic2payloads.put(topic.name(), schemaRandomPayloads(topic.name(), config.valueSchema, config.randomPoolSize, config.producerConfigs)); + } + }); + return topic2payloads::get; + } } /** @@ -189,13 +309,10 @@ private List randomPayloads(int size, double randomRatio, int count) { int randomBytes = (int) (size * randomRatio); int staticBytes = size - randomBytes; byte[] staticPayload = new byte[staticBytes]; - r.nextBytes(staticPayload); - if (randomBytes == 0) { // all payloads are the same, no need to create multiple copies return List.of(staticPayload); } - List payloads = new ArrayList<>(count); for (int i = 0; i < count; i++) { byte[] payload = new byte[size]; @@ -237,4 +354,133 @@ public void close() { producerService.close(); consumerService.close(); } + + private static List schemaPayloads(String topic, String schemaJson, String payloadsFile, Map configs) { + try (KafkaAvroSerializer serializer = new KafkaAvroSerializer()) { + List payloads = new ArrayList<>(); + AvroSchema schema = new AvroSchema(schemaJson); + serializer.configure(configs, false); + for (String payloadStr : Files.readAllLines(Path.of(payloadsFile), StandardCharsets.UTF_8)) { + Object object = AvroSchemaUtils.toObject(payloadStr, schema); + byte[] payload = serializer.serialize(topic, object); + payloads.add(payload); + } + return payloads; + } catch (IOException ex) { + throw new RuntimeException(ex); + } + } + + private static List schemaRandomPayloads(String topic, String schemaJson, int count, Map configs) { + try (KafkaAvroSerializer serializer = new KafkaAvroSerializer()) { + List payloads = new ArrayList<>(count); + AvroSchema avroSchema = new AvroSchema(schemaJson); + Schema schema = avroSchema.rawSchema(); + serializer.configure(configs, false); + for (int i = 0; i < count; i++) { + Object value = randomAvroValue(schema); + payloads.add(serializer.serialize(topic, value)); + } + return payloads; + } + } + + @SuppressWarnings("checkstyle:cyclomaticComplexity") + private static Object randomAvroValue(Schema schema) { + switch (schema.getType()) { + case NULL: + return null; + case BOOLEAN: + return ThreadLocalRandom.current().nextBoolean(); + case INT: + return ThreadLocalRandom.current().nextInt(); + case LONG: + return ThreadLocalRandom.current().nextLong(); + case FLOAT: + return ThreadLocalRandom.current().nextFloat(); + case DOUBLE: + return ThreadLocalRandom.current().nextDouble(); + case BYTES: { + int len = ThreadLocalRandom.current().nextInt(1, 17); + byte[] b = new byte[len]; + ThreadLocalRandom.current().nextBytes(b); + return ByteBuffer.wrap(b); + } + case STRING: + return randomString(ThreadLocalRandom.current().nextInt(3, 16)); + case ENUM: { + List symbols = schema.getEnumSymbols(); + return new GenericData.EnumSymbol(schema, symbols.get(ThreadLocalRandom.current().nextInt(symbols.size()))); + } + case FIXED: { + int size = schema.getFixedSize(); + byte[] b = new byte[size]; + ThreadLocalRandom.current().nextBytes(b); + return new GenericData.Fixed(schema, b); + } + case RECORD: { + GenericRecord rec = new GenericData.Record(schema); + for (Schema.Field f : schema.getFields()) { + Schema fSchema = effectiveSchema(f.schema()); + rec.put(f.name(), randomAvroValue(fSchema)); + } + return rec; + } + case ARRAY: { + int n = ThreadLocalRandom.current().nextInt(0, 4); + List arr = new ArrayList<>(n); + Schema elemSchema = effectiveSchema(schema.getElementType()); + for (int i = 0; i < n; i++) { + arr.add(randomAvroValue(elemSchema)); + } + return arr; + } + case MAP: { + int n = ThreadLocalRandom.current().nextInt(0, 4); + Map map = new HashMap<>(n); + Schema valSchema = effectiveSchema(schema.getValueType()); + for (int i = 0; i < n; i++) { + map.put("k" + i, randomAvroValue(valSchema).toString()); + } + return map; + } + case UNION: { + // Prefer a non-null type if present + List types = schema.getTypes(); + List nonNull = new ArrayList<>(); + for (Schema s : types) { + if (s.getType() != Schema.Type.NULL) nonNull.add(s); + } + Schema chosen; + if (!nonNull.isEmpty()) { + chosen = nonNull.get(ThreadLocalRandom.current().nextInt(nonNull.size())); + } else { + chosen = types.get(ThreadLocalRandom.current().nextInt(types.size())); + } + return randomAvroValue(chosen); + } + default: + // Fallback to string for unhandled types + return randomString(8); + } + } + + private static Schema effectiveSchema(Schema schema) { + if (schema.getType() == Schema.Type.UNION) { + for (Schema s : schema.getTypes()) { + if (s.getType() != Schema.Type.NULL) return s; + } + } + return schema; + } + + private static String randomString(int length) { + final String characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; + Random r = ThreadLocalRandom.current(); + StringBuilder sb = new StringBuilder(length); + for (int i = 0; i < length; i++) { + sb.append(characters.charAt(r.nextInt(characters.length()))); + } + return sb.toString(); + } } diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/ConsumerService.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/ConsumerService.java index faab997a98..c4b2c437bf 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/ConsumerService.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/ConsumerService.java @@ -1,17 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; import org.apache.kafka.clients.admin.Admin; +import org.apache.kafka.clients.admin.AdminClientConfig; import org.apache.kafka.clients.admin.AlterConsumerGroupOffsetsResult; import org.apache.kafka.clients.admin.ListOffsetsResult; import org.apache.kafka.clients.admin.OffsetSpec; @@ -21,14 +30,18 @@ import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.header.Header; import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.apache.kafka.common.serialization.StringDeserializer; import org.apache.kafka.common.utils.ThreadUtils; import org.apache.kafka.tools.automq.perf.TopicService.Topic; +import com.google.common.primitives.Longs; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -41,7 +54,7 @@ import java.util.Map; import java.util.Properties; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionStage; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -50,7 +63,9 @@ import java.util.stream.Collectors; import java.util.stream.Stream; -import static org.apache.kafka.tools.automq.perf.ProducerService.HEADER_KEY_CHARSET; +import io.github.bucket4j.BlockingBucket; +import io.github.bucket4j.Bucket; + import static org.apache.kafka.tools.automq.perf.ProducerService.HEADER_KEY_SEND_TIME_NANOS; public class ConsumerService implements AutoCloseable { @@ -61,11 +76,10 @@ public class ConsumerService implements AutoCloseable { private final List groups = new ArrayList<>(); private final String groupSuffix; - public ConsumerService(String bootstrapServer) { - Properties properties = new Properties(); - properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer); - properties.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, (int) TimeUnit.SECONDS.toMillis(300)); - properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); + public ConsumerService(String bootstrapServer, Properties properties) { + properties.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer); + properties.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, (int) TimeUnit.MINUTES.toMillis(2)); + properties.put(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, (int) TimeUnit.MINUTES.toMillis(10)); this.admin = Admin.create(properties); this.groupSuffix = new SimpleDateFormat("HHmmss").format(System.currentTimeMillis()); } @@ -79,20 +93,21 @@ public ConsumerService(String bootstrapServer) { * @param config consumer configuration * @return the number of consumers created */ - public int createConsumers(List topics, ConsumersConfig config) { + public int createConsumers(List topics, ConsumersConfig config, Stats stats) { int count = 0; for (int g = 0; g < config.groupsPerTopic; g++) { - Group group = new Group(g, config.consumersPerGroup, topics, config); + Group group = new Group(g, config.consumersPerGroup, topics, config, stats); groups.add(group); - count += group.size(); + count += group.consumerCount(); } return count; } - public void start(ConsumerCallback callback) { + public void start(ConsumerCallback callback, int pollRate) { + BlockingBucket bucket = rateLimitBucket(pollRate); CompletableFuture.allOf( groups.stream() - .map(group -> group.start(callback)) + .map(group -> group.start(callback, bucket)) .toArray(CompletableFuture[]::new) ).join(); } @@ -101,17 +116,111 @@ public void pause() { groups.forEach(Group::pause); } + /** + * Resume all consumer groups + */ public void resume() { groups.forEach(Group::resume); } + /** + * Resume only a percentage of consumer groups + * + * @param percentage The percentage of consumers to resume (0-100) + */ + public void resume(int percentage) { + int size = groups.size(); + int consumersToResume = (int) Math.ceil(size * (percentage / 100.0)); + consumersToResume = Math.max(1, Math.min(size, consumersToResume)); // Ensure at least 1 and at most size + + LOGGER.info("Resuming {}% of consumers ({} out of {})", percentage, consumersToResume, size); + + for (int i = 0; i < consumersToResume; i++) { + groups.get(i).resume(); + } + } + + /** + * Pause all consumers for the specified topics across all groups. + */ + public void pauseTopics(Collection topics) { + for (Group group : groups) { + for (Topic topic : group.consumers.keySet()) { + if (topics.contains(topic.name)) { + List topicConsumers = group.consumers.get(topic); + if (topicConsumers != null) { + topicConsumers.forEach(Consumer::pause); + } + } + } + } + } + + /** + * Resume all consumers for the specified topics across all groups. + */ + public void resumeTopics(Collection topics) { + for (Group group : groups) { + for (Topic topic : group.consumers.keySet()) { + if (topics.contains(topic.name)) { + List topicConsumers = group.consumers.get(topic); + if (topicConsumers != null) { + topicConsumers.forEach(Consumer::resume); + } + } + } + } + } + + /** + * Reset consumer offsets for catch-up reading. + * + * @param startMillis The timestamp to start seeking from + * @param intervalMillis The interval between group starts + * @param percentage The percentage of consumers to activate (0-100) + */ + public void resetOffset(long startMillis, long intervalMillis, int percentage) { + AtomicLong timestamp = new AtomicLong(startMillis); + int size = groups.size(); + int consumersToActivate = (int) Math.ceil(size * (percentage / 100.0)); + consumersToActivate = Math.max(1, Math.min(size, consumersToActivate)); // Ensure at least 1 and at most size + + LOGGER.info("Activating {}% of consumers ({} out of {})", percentage, consumersToActivate, size); + + for (int i = 0; i < consumersToActivate; i++) { + Group group = groups.get(i); + group.seek(timestamp.getAndAdd(intervalMillis)); + LOGGER.info("Reset consumer group offsets: {}/{}", i + 1, consumersToActivate); + } + + // Keep the remaining consumers paused + if (consumersToActivate < size) { + LOGGER.info("Keeping {} consumer groups paused during catch-up", size - consumersToActivate); + } + } + + /** + * Reset all consumer offsets (100% consumers) + * @param startMillis The timestamp to start seeking from + * @param intervalMillis The interval between group starts + */ public void resetOffset(long startMillis, long intervalMillis) { - AtomicLong start = new AtomicLong(startMillis); - CompletableFuture.allOf( - groups.stream() - .map(group -> group.seek(start.getAndAdd(intervalMillis))) - .toArray(CompletableFuture[]::new) - ).join(); + resetOffset(startMillis, intervalMillis, 100); + } + + public int consumerCount() { + return groups.stream() + .mapToInt(Group::consumerCount) + .sum(); + } + + private BlockingBucket rateLimitBucket(int rateLimit) { + return Bucket.builder() + .addLimit(limit -> limit + .capacity(rateLimit / 10) + .refillGreedy(rateLimit, Duration.ofSeconds(1)) + ).build() + .asBlocking(); } @Override @@ -125,24 +234,23 @@ public interface ConsumerCallback { /** * Called when a message is received. * - * @param payload the received message payload - * @param sendTimeNanos the time in nanoseconds when the message was sent + * @param topicPartition the topic partition of the received message */ - void messageReceived(byte[] payload, long sendTimeNanos); + void messageReceived(TopicPartition topicPartition) throws InterruptedException; } public static class ConsumersConfig { final String bootstrapServer; final int groupsPerTopic; final int consumersPerGroup; - final Map consumerConfigs; + final Properties properties; public ConsumersConfig(String bootstrapServer, int groupsPerTopic, int consumersPerGroup, - Map consumerConfigs) { + Properties properties) { this.bootstrapServer = bootstrapServer; this.groupsPerTopic = groupsPerTopic; this.consumersPerGroup = consumersPerGroup; - this.consumerConfigs = consumerConfigs; + this.properties = properties; } } @@ -150,22 +258,22 @@ private class Group implements AutoCloseable { private final int index; private final Map> consumers = new HashMap<>(); - public Group(int index, int consumersPerGroup, List topics, ConsumersConfig config) { + public Group(int index, int consumersPerGroup, List topics, ConsumersConfig config, Stats stats) { this.index = index; Properties common = toProperties(config); for (Topic topic : topics) { List topicConsumers = new ArrayList<>(); for (int c = 0; c < consumersPerGroup; c++) { - Consumer consumer = newConsumer(topic, common); + Consumer consumer = newConsumer(topic, common, stats); topicConsumers.add(consumer); } consumers.put(topic, topicConsumers); } } - public CompletableFuture start(ConsumerCallback callback) { - consumers().forEach(consumer -> consumer.start(callback)); + public CompletableFuture start(ConsumerCallback callback, BlockingBucket bucket) { + consumers().forEach(consumer -> consumer.start(callback, bucket)); // wait for all consumers to join the group return CompletableFuture.allOf(consumers() @@ -181,20 +289,27 @@ public void resume() { consumers().forEach(Consumer::resume); } - public CompletableFuture seek(long timestamp) { - return admin.listOffsets(listOffsetsRequest(timestamp)) - .all() - .toCompletionStage() - .toCompletableFuture() - .thenCompose(offsetMap -> CompletableFuture.allOf(consumers.keySet().stream() - .map(topic -> admin.alterConsumerGroupOffsets(groupId(topic), resetOffsetsRequest(topic, offsetMap))) + public void seek(long timestamp) { + // assuming all partitions approximately have the same offset at the given timestamp + TopicPartition firstPartition = consumers.keySet().iterator().next().firstPartition(); + try { + ListOffsetsResult.ListOffsetsResultInfo offsetInfo = admin.listOffsets(Map.of(firstPartition, OffsetSpec.forTimestamp(timestamp))) + .partitionResult(firstPartition) + .get(); + KafkaFuture.allOf(consumers.keySet().stream() + .map(topic -> admin.alterConsumerGroupOffsets(groupId(topic), resetOffsetsRequest(topic, offsetInfo.offset()))) .map(AlterConsumerGroupOffsetsResult::all) - .map(KafkaFuture::toCompletionStage) - .map(CompletionStage::toCompletableFuture) - .toArray(CompletableFuture[]::new))); + .toArray(KafkaFuture[]::new)).get(); + } catch (ExecutionException | InterruptedException e) { + throw new RuntimeException("Failed to list or reset consumer offsets", e); + } + } + + public int consumerGroupCount() { + return consumers.size(); } - public int size() { + public int consumerCount() { return consumers.values().stream() .mapToInt(List::size) .sum(); @@ -208,18 +323,19 @@ public void close() { private Properties toProperties(ConsumersConfig config) { Properties properties = new Properties(); - properties.putAll(config.consumerConfigs); + properties.putAll(config.properties); properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, config.bootstrapServer); properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName()); properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName()); + properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, OffsetResetStrategy.LATEST.toString()); return properties; } - private Consumer newConsumer(Topic topic, Properties common) { + private Consumer newConsumer(Topic topic, Properties common, Stats stats) { Properties properties = new Properties(); properties.putAll(common); properties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId(topic)); - return new Consumer(properties, topic.name); + return new Consumer(properties, topic.name, stats); } private Stream consumers() { @@ -227,26 +343,14 @@ private Stream consumers() { } private String groupId(Topic topic) { - return String.format("sub-%s-%s-%03d", topic.name, groupSuffix, index); + return String.format("sub-%s-%03d", topic.name, index); } - private Map listOffsetsRequest(long timestamp) { - return consumers.keySet().stream() - .map(Topic::partitions) - .flatMap(List::stream) + private Map resetOffsetsRequest(Topic topic, long offset) { + return topic.partitions().stream() .collect(Collectors.toMap( partition -> partition, - partition -> OffsetSpec.forTimestamp(timestamp) - )); - } - - private Map resetOffsetsRequest(Topic topic, - Map offsetMap) { - return offsetMap.entrySet().stream() - .filter(entry -> topic.containsPartition(entry.getKey())) - .collect(Collectors.toMap( - Map.Entry::getKey, - entry -> new OffsetAndMetadata(entry.getValue().offset()) + ignore -> new OffsetAndMetadata(offset) )); } } @@ -260,16 +364,17 @@ private static class Consumer { private final CompletableFuture started = new CompletableFuture<>(); private boolean paused = false; private volatile boolean closing = false; + private final Stats stats; - public Consumer(Properties properties, String topic) { + public Consumer(Properties properties, String topic, Stats stats) { this.consumer = new KafkaConsumer<>(properties); this.executor = Executors.newSingleThreadExecutor(ThreadUtils.createThreadFactory("perf-consumer", false)); - + this.stats = stats; consumer.subscribe(List.of(topic), subscribeListener()); } - public void start(ConsumerCallback callback) { - this.task = this.executor.submit(() -> pollRecords(consumer, callback)); + public void start(ConsumerCallback callback, BlockingBucket bucket) { + this.task = this.executor.submit(() -> pollRecords(consumer, callback, bucket)); } public CompletableFuture started() { @@ -299,18 +404,29 @@ public void onPartitionsAssigned(Collection partitions) { }; } - private void pollRecords(KafkaConsumer consumer, ConsumerCallback callback) { + private void pollRecords(KafkaConsumer consumer, ConsumerCallback callback, BlockingBucket bucket) { while (!closing) { try { while (paused) { Thread.sleep(PAUSE_INTERVAL); } ConsumerRecords records = consumer.poll(POLL_TIMEOUT); + int numMessages = records.count(); + if (numMessages == 0) { + continue; + } + ConsumerRecord firstRecord = records.iterator().next(); + Header header = firstRecord.headers().lastHeader(HEADER_KEY_SEND_TIME_NANOS); + long bytes = 0; + long sendTimeNanos = Longs.fromByteArray(header.value()); for (ConsumerRecord record : records) { - long sendTimeNanos = Long.parseLong(new String(record.headers().lastHeader(HEADER_KEY_SEND_TIME_NANOS).value(), HEADER_KEY_CHARSET)); - callback.messageReceived(record.value(), sendTimeNanos); + TopicPartition topicPartition = new TopicPartition(record.topic(), record.partition()); + bytes += record.value().length; + callback.messageReceived(topicPartition); } - } catch (InterruptException e) { + stats.messageReceived(numMessages, bytes, sendTimeNanos); + bucket.consume(records.count()); + } catch (InterruptException | InterruptedException e) { // ignore, as we are closing } catch (Exception e) { LOGGER.warn("exception occur while consuming message", e); diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/CpuMonitor.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/CpuMonitor.java new file mode 100644 index 0000000000..85f18f116d --- /dev/null +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/CpuMonitor.java @@ -0,0 +1,35 @@ +package org.apache.kafka.tools.automq.perf; + +import oshi.SystemInfo; +import oshi.hardware.CentralProcessor; + +public class CpuMonitor { + private final CentralProcessor processor; + private long[] prevTicks; + + public CpuMonitor() { + this.processor = new SystemInfo().getHardware().getProcessor(); + this.prevTicks = processor.getSystemCpuLoadTicks(); + } + + /** + * Returns the CPU usage between the last call of this method and now. + * It returns -1.0 if an error occurs. + * + * @return CPU load between 0 and 1 (100%) + */ + public synchronized double usage() { + try { + return usage0(); + } catch (Exception e) { + return -1.0; + } + } + + private double usage0() { + long[] currTicks = processor.getSystemCpuLoadTicks(); + double usage = processor.getSystemCpuLoadBetweenTicks(prevTicks); + prevTicks = currTicks; + return usage; + } +} diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/MemoryMonitor.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/MemoryMonitor.java new file mode 100644 index 0000000000..e876652755 --- /dev/null +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/MemoryMonitor.java @@ -0,0 +1,26 @@ +package org.apache.kafka.tools.automq.perf; + +import java.lang.management.BufferPoolMXBean; +import java.lang.management.ManagementFactory; +import java.util.List; + +public class MemoryMonitor { + + /** + * Returns the amount of heap memory used by the JVM in bytes. + */ + public static long heapUsed() { + return Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory(); + } + + /** + * Returns the amount of direct memory used by the JVM in bytes. + */ + public static long directUsed() { + List pools = ManagementFactory.getPlatformMXBeans(BufferPoolMXBean.class); + return pools.stream() + .filter(p -> "direct".equals(p.getName())) + .mapToLong(BufferPoolMXBean::getMemoryUsed) + .sum(); + } +} diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/PaddingDecimalFormat.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/PaddingDecimalFormat.java index a5b448417f..a133fc5190 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/PaddingDecimalFormat.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/PaddingDecimalFormat.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/PerfConfig.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/PerfConfig.java index f85160105a..33a87c7621 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/PerfConfig.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/PerfConfig.java @@ -1,17 +1,26 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; import org.apache.kafka.common.utils.Exit; +import org.apache.kafka.common.utils.Utils; import org.apache.kafka.tools.automq.perf.ConsumerService.ConsumersConfig; import org.apache.kafka.tools.automq.perf.ProducerService.ProducersConfig; import org.apache.kafka.tools.automq.perf.TopicService.TopicsConfig; @@ -24,17 +33,24 @@ import net.sourceforge.argparse4j.inf.Namespace; import net.sourceforge.argparse4j.internal.HelpScreenException; +import java.io.IOException; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.Properties; +import java.util.Random; +import java.util.concurrent.ThreadLocalRandom; import static net.sourceforge.argparse4j.impl.Arguments.storeTrue; +import static org.apache.kafka.tools.automq.perf.PerfConfig.IntegerArgumentType.between; import static org.apache.kafka.tools.automq.perf.PerfConfig.IntegerArgumentType.nonNegativeInteger; +import static org.apache.kafka.tools.automq.perf.PerfConfig.IntegerArgumentType.notLessThan; import static org.apache.kafka.tools.automq.perf.PerfConfig.IntegerArgumentType.positiveInteger; public class PerfConfig { public final String bootstrapServer; - public final Map commonConfigs; + public final Properties commonConfigs; public final Map topicConfigs; public final Map producerConfigs; public final Map consumerConfigs; @@ -45,16 +61,23 @@ public class PerfConfig { public final int producersPerTopic; public final int groupsPerTopic; public final int consumersPerGroup; + public final boolean awaitTopicReady; public final int recordSize; public final double randomRatio; public final int randomPoolSize; public final int sendRate; public final int sendRateDuringCatchup; + public final int maxConsumeRecordRate; public final int backlogDurationSeconds; public final int groupStartDelaySeconds; public final int warmupDurationMinutes; + public final int consumersDuringCatchupPercentage; public final int testDurationMinutes; public final int reportingIntervalSeconds; + public final String valueSchema; + public final String valuesFile; + public final boolean reuseTopics; // Added for --reuse-topics + public final String catchupTopicPrefix; // Added for --catchup-topic-prefix public PerfConfig(String[] args) { ArgumentParser parser = parser(); @@ -71,27 +94,36 @@ public PerfConfig(String[] args) { assert ns != null; bootstrapServer = ns.getString("bootstrapServer"); - commonConfigs = parseConfigs(ns.getList("commonConfigs")); + commonConfigs = ns.getString("commonConfigFile") == null ? new Properties() : loadProperties(ns.getString("commonConfigFile")); topicConfigs = parseConfigs(ns.getList("topicConfigs")); producerConfigs = parseConfigs(ns.getList("producerConfigs")); consumerConfigs = parseConfigs(ns.getList("consumerConfigs")); reset = ns.getBoolean("reset"); - topicPrefix = ns.getString("topicPrefix") == null ? "test-topic-" + System.currentTimeMillis() : ns.getString("topicPrefix"); + topicPrefix = ns.getString("topicPrefix") == null ? randomTopicPrefix() : ns.getString("topicPrefix"); topics = ns.getInt("topics"); partitionsPerTopic = ns.getInt("partitionsPerTopic"); producersPerTopic = ns.getInt("producersPerTopic"); groupsPerTopic = ns.getInt("groupsPerTopic"); consumersPerGroup = ns.getInt("consumersPerGroup"); + awaitTopicReady = ns.getBoolean("awaitTopicReady"); recordSize = ns.getInt("recordSize"); randomRatio = ns.getDouble("randomRatio"); randomPoolSize = ns.getInt("randomPoolSize"); - sendRate = ns.getInt("sendRate"); + sendRate = Optional.ofNullable(ns.getDouble("sendThroughput")) + .map(throughput -> (int) (throughput * 1024 * 1024 / recordSize)) + .orElse(ns.getInt("sendRate")); sendRateDuringCatchup = ns.getInt("sendRateDuringCatchup") == null ? sendRate : ns.getInt("sendRateDuringCatchup"); + maxConsumeRecordRate = ns.getInt("maxConsumeRecordRate"); backlogDurationSeconds = ns.getInt("backlogDurationSeconds"); groupStartDelaySeconds = ns.getInt("groupStartDelaySeconds"); warmupDurationMinutes = ns.getInt("warmupDurationMinutes"); testDurationMinutes = ns.getInt("testDurationMinutes"); reportingIntervalSeconds = ns.getInt("reportingIntervalSeconds"); + valueSchema = ns.getString("valueSchema"); + valuesFile = ns.get("valuesFile"); + reuseTopics = ns.getBoolean("reuseTopics"); // Added for --reuse-topics + catchupTopicPrefix = ns.getString("catchupTopicPrefix"); // Added for --catchup-topic-prefix + consumersDuringCatchupPercentage = ns.getInt("consumersDuringCatchupPercentage"); if (backlogDurationSeconds < groupsPerTopic * groupStartDelaySeconds) { throw new IllegalArgumentException(String.format("BACKLOG_DURATION_SECONDS(%d) should not be less than GROUPS_PER_TOPIC(%d) * GROUP_START_DELAY_SECONDS(%d)", @@ -104,18 +136,40 @@ public static ArgumentParser parser() { .newArgumentParser("performance-test") .defaultHelp(true) .description("This tool is used to run performance tests."); + + addConnectionArguments(parser); + addTopicArguments(parser); + addConsumerArguments(parser); + addProducerArguments(parser); + addTestConfigArguments(parser); + return parser; + } + + private static void addConnectionArguments(ArgumentParser parser) { parser.addArgument("-B", "--bootstrap-server") .setDefault("localhost:9092") .type(String.class) .dest("bootstrapServer") .metavar("BOOTSTRAP_SERVER") .help("The AutoMQ bootstrap server."); - parser.addArgument("-A", "--common-configs") - .nargs("*") + + parser.addArgument("--reuse-topics") + .action(storeTrue()) + .dest("reuseTopics") + .help("Reuse existing topics with the given prefix instead of creating new ones."); + + // Add the new parameter + parser.addArgument("--catchup-topic-prefix") + .type(String.class) + .dest("catchupTopicPrefix") + .metavar("CATCHUP_TOPIC_PREFIX") + .help("The topic prefix for catch-up read testing. Reuses existing topics with this prefix and skips message accumulation phase."); + + parser.addArgument("-F", "--common-config-file") .type(String.class) - .dest("commonConfigs") - .metavar("COMMON_CONFIG") - .help("The common configurations."); + .dest("commonConfigFile") + .metavar("COMMON_CONFIG_FILE") + .help("The property file containing common configurations to be passed to all clients —— producer, consumer, and admin."); parser.addArgument("-T", "--topic-configs") .nargs("*") .type(String.class) @@ -138,6 +192,9 @@ public static ArgumentParser parser() { .action(storeTrue()) .dest("reset") .help("delete all topics before running the test."); + } + + private static void addTopicArguments(ArgumentParser parser) { parser.addArgument("-X", "--topic-prefix") .type(String.class) .dest("topicPrefix") @@ -155,12 +212,9 @@ public static ArgumentParser parser() { .dest("partitionsPerTopic") .metavar("PARTITIONS_PER_TOPIC") .help("The number of partitions per topic."); - parser.addArgument("-p", "--producers-per-topic") - .setDefault(1) - .type(positiveInteger()) - .dest("producersPerTopic") - .metavar("PRODUCERS_PER_TOPIC") - .help("The number of producers per topic."); + } + + private static void addConsumerArguments(ArgumentParser parser) { parser.addArgument("-g", "--groups-per-topic") .setDefault(1) .type(nonNegativeInteger()) @@ -173,6 +227,34 @@ public static ArgumentParser parser() { .dest("consumersPerGroup") .metavar("CONSUMERS_PER_GROUP") .help("The number of consumers per group."); + parser.addArgument("--await-topic-ready") + .setDefault(true) + .type(Boolean.class) + .dest("awaitTopicReady") + .metavar("AWAIT_TOPIC_READY") + .help("Use produce / consume detect to check topic readiness."); + parser.addArgument("-m", "--max-consume-record-rate") + .setDefault(1_000_000_000) + .type(between(0, 1_000_000_000)) + .dest("maxConsumeRecordRate") + .metavar("MAX_CONSUME_RECORD_RATE") + .help("The max rate of consuming records per second."); + + parser.addArgument("--consumers-during-catchup") + .setDefault(100) + .type(between(0, 100)) + .dest("consumersDuringCatchupPercentage") + .metavar("CONSUMERS_DURING_CATCHUP_PERCENTAGE") + .help("The percentage of consumers to activate during catch-up read (0-100). Default is 100 (all consumers)."); + } + + private static void addProducerArguments(ArgumentParser parser) { + parser.addArgument("-p", "--producers-per-topic") + .setDefault(1) + .type(positiveInteger()) + .dest("producersPerTopic") + .metavar("PRODUCERS_PER_TOPIC") + .help("The number of producers per topic."); parser.addArgument("-s", "--record-size") .setDefault(1024) .type(positiveInteger()) @@ -197,14 +279,22 @@ public static ArgumentParser parser() { .dest("sendRate") .metavar("SEND_RATE") .help("The send rate in messages per second."); + parser.addArgument("-f", "--send-throughput") + .type(Double.class) + .dest("sendThroughput") + .metavar("SEND_THROUGHPUT") + .help("The send throughput in MB/s. If not set, the send rate will be used. This is an alternative to --send-rate."); parser.addArgument("-a", "--send-rate-during-catchup") .type(positiveInteger()) .dest("sendRateDuringCatchup") .metavar("SEND_RATE_DURING_CATCHUP") .help("The send rate in messages per second during catchup. If not set, the send rate will be used."); + } + + private static void addTestConfigArguments(ArgumentParser parser) { parser.addArgument("-b", "--backlog-duration") .setDefault(0) - .type(nonNegativeInteger()) + .type(notLessThan(300)) .dest("backlogDurationSeconds") .metavar("BACKLOG_DURATION_SECONDS") .help("The backlog duration in seconds, and zero means no backlog. Should not be less than GROUPS_PER_TOPIC * GROUP_START_DELAY_SECONDS."); @@ -232,20 +322,30 @@ public static ArgumentParser parser() { .dest("reportingIntervalSeconds") .metavar("REPORTING_INTERVAL_SECONDS") .help("The reporting interval in seconds."); - return parser; + parser.addArgument("--value-schema") + .type(String.class) + .dest("valueSchema") + .metavar("VALUE_SCHEMA") + .help("The schema of the values ex. {\"type\":\"record\",\"name\":\"myrecord\",\"fields\":[{\"name\":\"f1\",\"type\":\"string\"}]}"); + parser.addArgument("--values-file") + .type(String.class) + .dest("valuesFile") + .metavar("VALUES_FILE") + .help("The avro value file. Example file content {\"f1\": \"value1\"}"); + // removed --values-random; default to random when no --values-file is provided } public String bootstrapServer() { return bootstrapServer; } - public Map adminConfig() { - return commonConfigs; + public Properties adminConfig() { + Properties properties = new Properties(); + properties.putAll(commonConfigs); + return properties; } public TopicsConfig topicsConfig() { - Map topicConfigs = new HashMap<>(commonConfigs); - topicConfigs.putAll(this.topicConfigs); return new TopicsConfig( topicPrefix, topics, @@ -255,26 +355,36 @@ public TopicsConfig topicsConfig() { } public ProducersConfig producersConfig() { - Map producerConfigs = new HashMap<>(commonConfigs); - producerConfigs.putAll(this.producerConfigs); + Properties properties = new Properties(); + properties.putAll(commonConfigs); + properties.putAll(producerConfigs); return new ProducersConfig( bootstrapServer, producersPerTopic, - producerConfigs + properties ); } public ConsumersConfig consumersConfig() { - Map consumerConfigs = new HashMap<>(commonConfigs); - consumerConfigs.putAll(this.consumerConfigs); + Properties properties = new Properties(); + properties.putAll(commonConfigs); + properties.putAll(consumerConfigs); return new ConsumersConfig( bootstrapServer, groupsPerTopic, consumersPerGroup, - consumerConfigs + properties ); } + private Properties loadProperties(String filename) { + try { + return Utils.loadProps(filename); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + private Map parseConfigs(List configs) { if (configs == null) { return new HashMap<>(); @@ -290,6 +400,20 @@ private Map parseConfigs(List configs) { return map; } + private String randomTopicPrefix() { + return String.format("topic_%d_%s", System.currentTimeMillis(), randomString(4)); + } + + private String randomString(int length) { + final String characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; + Random r = ThreadLocalRandom.current(); + StringBuilder sb = new StringBuilder(length); + for (int i = 0; i < length; i++) { + sb.append(characters.charAt(r.nextInt(characters.length()))); + } + return sb.toString(); + } + static class IntegerArgumentType extends ReflectArgumentType { private final IntegerValidator validator; @@ -316,6 +440,14 @@ public static IntegerArgumentType nonNegativeInteger() { public static IntegerArgumentType positiveInteger() { return new IntegerArgumentType(value -> value <= 0 ? "expected a positive integer, but got " + value : null); } + + public static IntegerArgumentType notLessThan(int min) { + return new IntegerArgumentType(value -> value < min ? "expected an integer not less than " + min + ", but got " + value : null); + } + + public static IntegerArgumentType between(int min, int max) { + return new IntegerArgumentType(value -> value < min || value > max ? "expected an integer between " + min + " and " + max + ", but got " + value : null); + } } @FunctionalInterface diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/ProducerService.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/ProducerService.java index cc6805200c..87d2a3b934 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/ProducerService.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/ProducerService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; @@ -22,6 +30,9 @@ import org.apache.kafka.common.utils.ThreadUtils; import org.apache.kafka.tools.automq.perf.TopicService.Topic; +import com.automq.stream.utils.Threads; +import com.google.common.primitives.Longs; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -41,6 +52,7 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -49,13 +61,14 @@ public class ProducerService implements AutoCloseable { public static final Charset HEADER_KEY_CHARSET = StandardCharsets.UTF_8; - public static final String HEADER_KEY_SEND_TIME_NANOS = "send_time_nanos"; + public static final String HEADER_KEY_SEND_TIME_NANOS = "st"; public static final double MIN_RATE = 1.0; private static final int ADJUST_RATE_INTERVAL_SECONDS = 5; private static final Logger LOGGER = LoggerFactory.getLogger(ProducerService.class); private final List producers = new LinkedList<>(); - private final ScheduledExecutorService adjustRateExecutor = Executors.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory("perf-producer-rate-adjust", true)); + private final ScheduledExecutorService adjustRateExecutor = + Threads.newSingleThreadScheduledExecutor("perf-producer-rate-adjust", true, LOGGER); private final ExecutorService executor = Executors.newCachedThreadPool(ThreadUtils.createThreadFactory("perf-producer", false)); /** @@ -92,10 +105,11 @@ public int createProducers(List topics, ProducersConfig config, ProducerC private Producer createProducer(Topic topic, ProducersConfig config, ProducerCallback callback) { Properties properties = new Properties(); - properties.putAll(config.producerConfigs); + properties.putAll(config.properties); properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, config.bootstrapServer); properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName()); properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName()); + properties.put(ProducerConfig.PARTITIONER_IGNORE_KEYS_CONFIG, true); KafkaProducer kafkaProducer = new KafkaProducer<>(properties); return new Producer(kafkaProducer, topic, callback); @@ -120,7 +134,7 @@ public int probe() { /** * Start sending messages using the given payloads at the given rate. */ - public void start(List payloads, double rate) { + public void start(Function> payloads, double rate) { adjustRate(rate); adjustRateExecutor.scheduleWithFixedDelay(this::adjustRate, 0, ADJUST_RATE_INTERVAL_SECONDS, TimeUnit.SECONDS); int processors = Runtime.getRuntime().availableProcessors(); @@ -176,7 +190,7 @@ private double calculateY(long x1, double y1, long x2, double y2, long x) { return y1 + (x - x1) * (y2 - y1) / (x2 - x1); } - private void start(List producers, List payloads) { + private void start(List producers, Function> payloads) { executor.submit(() -> { try { sendMessages(producers, payloads); @@ -186,11 +200,14 @@ private void start(List producers, List payloads) { }); } - private void sendMessages(List producers, List payloads) { + private void sendMessages(List producers, Function> payloadsGet) { Random random = ThreadLocalRandom.current(); while (!closed) { producers.forEach( - p -> sendMessage(p, payloads.get(random.nextInt(payloads.size()))) + p -> { + List payloads = payloadsGet.apply(p.topic.name); + sendMessage(p, payloads.get(random.nextInt(payloads.size()))); + } ); } } @@ -234,12 +251,12 @@ public interface ProducerCallback { public static class ProducersConfig { final String bootstrapServer; final int producersPerTopic; - final Map producerConfigs; + final Properties properties; - public ProducersConfig(String bootstrapServer, int producersPerTopic, Map producerConfigs) { + public ProducersConfig(String bootstrapServer, int producersPerTopic, Properties properties) { this.bootstrapServer = bootstrapServer; this.producersPerTopic = producersPerTopic; - this.producerConfigs = producerConfigs; + this.properties = properties; } } @@ -261,8 +278,6 @@ private static class Producer implements AutoCloseable { private final Topic topic; private final ProducerCallback callback; - private int partitionIndex = 0; - public Producer(KafkaProducer producer, Topic topic, ProducerCallback callback) { this.producer = producer; this.topic = topic; @@ -274,11 +289,7 @@ public Producer(KafkaProducer producer, Topic topic, ProducerCal * NOT thread-safe. */ public CompletableFuture sendAsync(byte[] payload) { - return sendAsync(nextKey(), payload, nextPartition()); - } - - private int nextPartition() { - return partitionIndex++ % topic.partitions; + return sendAsync(nextKey(), payload, null); } private String nextKey() { @@ -290,7 +301,7 @@ private String nextKey() { */ public List> probe() { return IntStream.range(0, topic.partitions) - .mapToObj(i -> sendAsync("probe", new byte[42], i)) + .mapToObj(i -> sendAsync("probe", new byte[] {1}, i)) .collect(Collectors.toList()); } @@ -303,9 +314,9 @@ public List> probe() { * @return a future that completes when the message is sent */ private CompletableFuture sendAsync(String key, byte[] payload, Integer partition) { - long sendTimeNanos = System.nanoTime(); + long sendTimeNanos = StatsCollector.currentNanos(); List
    headers = List.of( - new RecordHeader(HEADER_KEY_SEND_TIME_NANOS, Long.toString(sendTimeNanos).getBytes(HEADER_KEY_CHARSET)) + new RecordHeader(HEADER_KEY_SEND_TIME_NANOS, Longs.toByteArray(sendTimeNanos)) ); ProducerRecord record = new ProducerRecord<>(topic.name, partition, key, payload, headers); int size = payload.length; diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/Stats.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/Stats.java index 1364764627..f0058504c0 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/Stats.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/Stats.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; @@ -48,7 +56,7 @@ public class Stats { public final AtomicLong maxSendTimeNanos = new AtomicLong(0); public void messageSent(long bytes, long sendTimeNanos) { - long latencyMicros = TimeUnit.NANOSECONDS.toMicros(System.nanoTime() - sendTimeNanos); + long latencyMicros = TimeUnit.NANOSECONDS.toMicros(StatsCollector.currentNanos() - sendTimeNanos); messagesSent.increment(); bytesSent.add(bytes); sendLatencyMicros.recordValue(latencyMicros); @@ -62,12 +70,12 @@ public void messageFailed() { totalMessagesSendFailed.increment(); } - public void messageReceived(long bytes, long sendTimeNanos) { - long latencyMicros = TimeUnit.NANOSECONDS.toMicros(System.nanoTime() - sendTimeNanos); - messagesReceived.increment(); + public void messageReceived(long numMessages, long bytes, long sendTimeNanos) { + long latencyMicros = TimeUnit.NANOSECONDS.toMicros(StatsCollector.currentNanos() - sendTimeNanos); + messagesReceived.add(numMessages); bytesReceived.add(bytes); endToEndLatencyMicros.recordValue(latencyMicros); - totalMessagesReceived.increment(); + totalMessagesReceived.add(numMessages); totalBytesReceived.add(bytes); totalEndToEndLatencyMicros.recordValue(latencyMicros); maxSendTimeNanos.updateAndGet(current -> Math.max(current, sendTimeNanos)); diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/StatsCollector.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/StatsCollector.java index 16141d8156..a79a85d121 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/StatsCollector.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/StatsCollector.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; @@ -18,6 +26,7 @@ import org.slf4j.LoggerFactory; import java.text.DecimalFormat; +import java.time.Instant; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -34,32 +43,39 @@ public class StatsCollector { private static final double BYTES_PER_MB = 1 << 20; private static final double BYTES_PER_GB = 1 << 30; - // max to 9999.9 (167 min) + // max to 9999.9 s (167 min) private static final DecimalFormat DURATION_FORMAT = new PaddingDecimalFormat("0.0", 6); - // max to 999999.99 (1M msg/s) + // max to 99.99% + private static final DecimalFormat PERCENT_FORMAT = new PaddingDecimalFormat("0.00", 5); + // max to 99999 MiB (100 GiB) + private static final DecimalFormat MEMORY_FORMAT = new PaddingDecimalFormat("0", 5); + // max to 999999.99 msg/s (1M msg/s) private static final DecimalFormat RATE_FORMAT = new PaddingDecimalFormat("0.00", 9); - // max to 999.99 (1 GB/s) + // max to 999.99 MiB/s (1 GiB/s) private static final DecimalFormat THROUGHPUT_FORMAT = new PaddingDecimalFormat("0.00", 6); - // max to 99999.999 (100 s) + // max to 99999.999 ms (100 s) private static final DecimalFormat LATENCY_FORMAT = new PaddingDecimalFormat("0.000", 9); + // max to 999.99 private static final DecimalFormat COUNT_FORMAT = new PaddingDecimalFormat("0.00", 6); private static final String PERIOD_LOG_FORMAT = "{}s" + + " | CPU {}% | Mem {} MiB heap / {} MiB direct" + " | Prod rate {} msg/s / {} MiB/s | Prod err {} err/s" + " | Cons rate {} msg/s / {} MiB/s | Backlog: {} K msg" + - " | Prod Latency (ms) avg: {} - 50%: {} - 99%: {} - 99.9%: {} - Max: {}" + - " | E2E Latency (ms) avg: {} - 50%: {} - 99%: {} - 99.9%: {} - Max: {}"; + " | Prod Latency (ms) avg: {} - min: {} - 50%: {} - 99%: {} - 99.9%: {} - max: {}" + + " | E2E Latency (ms) avg: {} - min: {} - 50%: {} - 99%: {} - 99.9%: {} - max: {}"; private static final String SUMMARY_LOG_FORMAT = "Summary" + " | Prod rate {} msg/s / {} MiB/s | Prod total {} M msg / {} GiB / {} K err" + " | Cons rate {} msg/s / {} MiB/s | Cons total {} M msg / {} GiB" + - " | Prod Latency (ms) avg: {} - 50%: {} - 75%: {} - 90%: {} - 95%: {} - 99%: {} - 99.9%: {} - 99.99%: {} - Max: {}" + - " | E2E Latency (ms) avg: {} - 50%: {} - 75%: {} - 90%: {} - 95%: {} - 99%: {} - 99.9%: {} - 99.99%: {} - Max: {}"; + " | Prod Latency (ms) avg: {} - min: {} - 50%: {} - 75%: {} - 90%: {} - 95%: {} - 99%: {} - 99.9%: {} - 99.99%: {} - max: {}" + + " | E2E Latency (ms) avg: {} - min: {} - 50%: {} - 75%: {} - 90%: {} - 95%: {} - 99%: {} - 99.9%: {} - 99.99%: {} - max: {}"; private static final Logger LOGGER = LoggerFactory.getLogger(StatsCollector.class); public static Result printAndCollectStats(Stats stats, StopCondition condition, long intervalNanos, PerfConfig config) { final long start = System.nanoTime(); + CpuMonitor cpu = new CpuMonitor(); Result result = new Result(config); long last = start; @@ -74,7 +90,7 @@ public static Result printAndCollectStats(Stats stats, StopCondition condition, double elapsed = (periodStats.nowNanos - last) / NANOS_PER_SEC; double elapsedTotal = (periodStats.nowNanos - start) / NANOS_PER_SEC; - PeriodResult periodResult = new PeriodResult(periodStats, elapsed, config.groupsPerTopic); + PeriodResult periodResult = new PeriodResult(cpu, periodStats, elapsed, config.groupsPerTopic); result.update(periodResult, elapsedTotal); periodResult.logIt(elapsedTotal); @@ -105,6 +121,11 @@ public interface StopCondition { boolean shouldStop(long startNanos, long nowNanos); } + public static long currentNanos() { + Instant currentTime = Instant.now(); + return TimeUnit.SECONDS.toNanos(currentTime.getEpochSecond()) + currentTime.getNano(); + } + @SuppressFBWarnings("URF_UNREAD_PUBLIC_OR_PROTECTED_FIELD") public static class Result { public final PerfConfig config; @@ -144,6 +165,7 @@ public static class Result { public final List consumeThroughputBps = new ArrayList<>(); public final List backlog = new ArrayList<>(); public final List produceLatencyMeanMicros = new ArrayList<>(); + public final List produceLatencyMinMicros = new ArrayList<>(); public final List produceLatency50thMicros = new ArrayList<>(); public final List produceLatency75thMicros = new ArrayList<>(); public final List produceLatency90thMicros = new ArrayList<>(); @@ -154,6 +176,7 @@ public static class Result { public final List produceLatencyMaxMicros = new ArrayList<>(); public Map produceLatencyQuantilesMicros; public final List endToEndLatencyMeanMicros = new ArrayList<>(); + public final List endToEndLatencyMinMicros = new ArrayList<>(); public final List endToEndLatency50thMicros = new ArrayList<>(); public final List endToEndLatency75thMicros = new ArrayList<>(); public final List endToEndLatency90thMicros = new ArrayList<>(); @@ -178,6 +201,7 @@ private void update(PeriodResult periodResult, double elapsedTotal) { this.consumeThroughputBps.add(periodResult.consumeThroughputBps); this.backlog.add(periodResult.backlog); this.produceLatencyMeanMicros.add(periodResult.produceLatencyMeanMicros); + this.produceLatencyMinMicros.add(periodResult.produceLatencyMinMicros); this.produceLatency50thMicros.add(periodResult.produceLatency50thMicros); this.produceLatency75thMicros.add(periodResult.produceLatency75thMicros); this.produceLatency90thMicros.add(periodResult.produceLatency90thMicros); @@ -187,6 +211,7 @@ private void update(PeriodResult periodResult, double elapsedTotal) { this.produceLatency9999thMicros.add(periodResult.produceLatency9999thMicros); this.produceLatencyMaxMicros.add(periodResult.produceLatencyMaxMicros); this.endToEndLatencyMeanMicros.add(periodResult.endToEndLatencyMeanMicros); + this.endToEndLatencyMinMicros.add(periodResult.endToEndLatencyMinMicros); this.endToEndLatency50thMicros.add(periodResult.endToEndLatency50thMicros); this.endToEndLatency75thMicros.add(periodResult.endToEndLatency75thMicros); this.endToEndLatency90thMicros.add(periodResult.endToEndLatency90thMicros); @@ -231,6 +256,9 @@ private void update(CumulativeResult cumulativeResult) { } private static class PeriodResult { + private final double cpuUsage; + private final long heapMemoryUsed; + private final long directMemoryUsed; private final double produceRate; private final double produceThroughputBps; private final double errorRate; @@ -238,6 +266,7 @@ private static class PeriodResult { private final double consumeThroughputBps; private final long backlog; private final double produceLatencyMeanMicros; + private final double produceLatencyMinMicros; private final double produceLatency50thMicros; private final double produceLatency75thMicros; private final double produceLatency90thMicros; @@ -247,6 +276,7 @@ private static class PeriodResult { private final double produceLatency9999thMicros; private final double produceLatencyMaxMicros; private final double endToEndLatencyMeanMicros; + private final double endToEndLatencyMinMicros; private final double endToEndLatency50thMicros; private final double endToEndLatency75thMicros; private final double endToEndLatency90thMicros; @@ -256,7 +286,10 @@ private static class PeriodResult { private final double endToEndLatency9999thMicros; private final double endToEndLatencyMaxMicros; - private PeriodResult(PeriodStats stats, double elapsed, int readWriteRatio) { + private PeriodResult(CpuMonitor cpu, PeriodStats stats, double elapsed, int readWriteRatio) { + this.cpuUsage = cpu.usage(); + this.heapMemoryUsed = MemoryMonitor.heapUsed(); + this.directMemoryUsed = MemoryMonitor.directUsed(); this.produceRate = stats.messagesSent / elapsed; this.produceThroughputBps = stats.bytesSent / elapsed; this.errorRate = stats.messagesSendFailed / elapsed; @@ -264,6 +297,7 @@ private PeriodResult(PeriodStats stats, double elapsed, int readWriteRatio) { this.consumeThroughputBps = stats.bytesReceived / elapsed; this.backlog = Math.max(0, readWriteRatio * stats.totalMessagesSent - stats.totalMessagesReceived); this.produceLatencyMeanMicros = stats.sendLatencyMicros.getMean(); + this.produceLatencyMinMicros = stats.sendLatencyMicros.getMinValue(); this.produceLatency50thMicros = stats.sendLatencyMicros.getValueAtPercentile(50); this.produceLatency75thMicros = stats.sendLatencyMicros.getValueAtPercentile(75); this.produceLatency90thMicros = stats.sendLatencyMicros.getValueAtPercentile(90); @@ -273,6 +307,7 @@ private PeriodResult(PeriodStats stats, double elapsed, int readWriteRatio) { this.produceLatency9999thMicros = stats.sendLatencyMicros.getValueAtPercentile(99.99); this.produceLatencyMaxMicros = stats.sendLatencyMicros.getMaxValue(); this.endToEndLatencyMeanMicros = stats.endToEndLatencyMicros.getMean(); + this.endToEndLatencyMinMicros = stats.endToEndLatencyMicros.getMinValue(); this.endToEndLatency50thMicros = stats.endToEndLatencyMicros.getValueAtPercentile(50); this.endToEndLatency75thMicros = stats.endToEndLatencyMicros.getValueAtPercentile(75); this.endToEndLatency90thMicros = stats.endToEndLatencyMicros.getValueAtPercentile(90); @@ -286,6 +321,9 @@ private PeriodResult(PeriodStats stats, double elapsed, int readWriteRatio) { private void logIt(double elapsedTotal) { LOGGER.info(PERIOD_LOG_FORMAT, DURATION_FORMAT.format(elapsedTotal), + PERCENT_FORMAT.format(cpuUsage * 100), + MEMORY_FORMAT.format(heapMemoryUsed / BYTES_PER_MB), + MEMORY_FORMAT.format(directMemoryUsed / BYTES_PER_MB), RATE_FORMAT.format(produceRate), THROUGHPUT_FORMAT.format(produceThroughputBps / BYTES_PER_MB), RATE_FORMAT.format(errorRate), @@ -293,11 +331,13 @@ private void logIt(double elapsedTotal) { THROUGHPUT_FORMAT.format(consumeThroughputBps / BYTES_PER_MB), COUNT_FORMAT.format(backlog / 1_000.0), LATENCY_FORMAT.format(produceLatencyMeanMicros / MICROS_PER_MILLI), + LATENCY_FORMAT.format(produceLatencyMinMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatency50thMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatency99thMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatency999thMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatencyMaxMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatencyMeanMicros / MICROS_PER_MILLI), + LATENCY_FORMAT.format(endToEndLatencyMinMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatency50thMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatency99thMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatency999thMicros / MICROS_PER_MILLI), @@ -317,6 +357,7 @@ private static class CumulativeResult { private final double consumeCountTotal; private final double consumeSizeTotalBytes; private final double produceLatencyMeanTotalMicros; + private final double produceLatencyMinTotalMicros; private final double produceLatency50thTotalMicros; private final double produceLatency75thTotalMicros; private final double produceLatency90thTotalMicros; @@ -327,6 +368,7 @@ private static class CumulativeResult { private final double produceLatencyMaxTotalMicros; public final Map produceLatencyQuantilesMicros = new TreeMap<>(); private final double endToEndLatencyMeanTotalMicros; + private final double endToEndLatencyMinTotalMicros; private final double endToEndLatency50thTotalMicros; private final double endToEndLatency75thTotalMicros; private final double endToEndLatency90thTotalMicros; @@ -348,6 +390,7 @@ private CumulativeResult(CumulativeStats stats, double elapsedTotal) { consumeCountTotal = stats.totalMessagesReceived; consumeSizeTotalBytes = stats.totalBytesReceived; produceLatencyMeanTotalMicros = stats.totalSendLatencyMicros.getMean(); + produceLatencyMinTotalMicros = stats.totalSendLatencyMicros.getMinValue(); produceLatency50thTotalMicros = stats.totalSendLatencyMicros.getValueAtPercentile(50); produceLatency75thTotalMicros = stats.totalSendLatencyMicros.getValueAtPercentile(75); produceLatency90thTotalMicros = stats.totalSendLatencyMicros.getValueAtPercentile(90); @@ -360,6 +403,7 @@ private CumulativeResult(CumulativeStats stats, double elapsedTotal) { value -> produceLatencyQuantilesMicros.put(value.getPercentile(), value.getValueIteratedTo()) ); endToEndLatencyMeanTotalMicros = stats.totalEndToEndLatencyMicros.getMean(); + endToEndLatencyMinTotalMicros = stats.totalEndToEndLatencyMicros.getMinValue(); endToEndLatency50thTotalMicros = stats.totalEndToEndLatencyMicros.getValueAtPercentile(50); endToEndLatency75thTotalMicros = stats.totalEndToEndLatencyMicros.getValueAtPercentile(75); endToEndLatency90thTotalMicros = stats.totalEndToEndLatencyMicros.getValueAtPercentile(90); @@ -385,6 +429,7 @@ private void logIt() { COUNT_FORMAT.format(consumeCountTotal / 1_000_000.0), COUNT_FORMAT.format(consumeSizeTotalBytes / BYTES_PER_GB), LATENCY_FORMAT.format(produceLatencyMeanTotalMicros / MICROS_PER_MILLI), + LATENCY_FORMAT.format(produceLatencyMinTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatency50thTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatency75thTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatency90thTotalMicros / MICROS_PER_MILLI), @@ -394,6 +439,7 @@ private void logIt() { LATENCY_FORMAT.format(produceLatency9999thTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(produceLatencyMaxTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatencyMeanTotalMicros / MICROS_PER_MILLI), + LATENCY_FORMAT.format(endToEndLatencyMinTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatency50thTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatency75thTotalMicros / MICROS_PER_MILLI), LATENCY_FORMAT.format(endToEndLatency90thTotalMicros / MICROS_PER_MILLI), diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/TopicService.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/TopicService.java index aad521d01d..ec53f91d52 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/TopicService.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/TopicService.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf; @@ -14,20 +22,25 @@ import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.AdminClientConfig; import org.apache.kafka.clients.admin.CreateTopicsResult; +import org.apache.kafka.clients.admin.DescribeTopicsResult; import org.apache.kafka.clients.admin.ListTopicsResult; import org.apache.kafka.clients.admin.NewTopic; +import org.apache.kafka.clients.admin.TopicDescription; +import org.apache.kafka.clients.admin.TopicListing; import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.errors.TopicExistsException; +import com.google.common.collect.Lists; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.HashMap; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Set; +import java.util.Properties; import java.util.concurrent.ExecutionException; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -35,12 +48,22 @@ public class TopicService implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(TopicService.class); + /** + * The maximum number of partitions per batch. + * + * @see org.apache.kafka.controller.ReplicationControlManager + */ + private static final int MAX_PARTITIONS_PER_BATCH = 10_000; + /** + * The common prefix for performance test topics. + */ + private static final String COMMON_TOPIC_PREFIX = "__automq_perf_"; + private final Admin admin; - public TopicService(String bootstrapServer, Map adminConfigs) { - Map configs = new HashMap<>(adminConfigs); - configs.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer); - this.admin = Admin.create(configs); + public TopicService(String bootstrapServer, Properties properties) { + properties.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer); + this.admin = Admin.create(properties); } /** @@ -49,26 +72,36 @@ public TopicService(String bootstrapServer, Map adminConfigs) { */ public List createTopics(TopicsConfig config) { List newTopics = IntStream.range(0, config.topics) - .mapToObj(i -> generateTopicName(config.topicPrefix, config.partitionsPerTopic, i)) - .map(name -> new NewTopic(name, config.partitionsPerTopic, (short) 1).configs(config.topicConfigs)) - .collect(Collectors.toList()); - CreateTopicsResult topics = admin.createTopics(newTopics); - topics.values().forEach(this::waitTopicCreated); - return topics.values().keySet().stream() - .map(name -> new Topic(name, config.partitionsPerTopic)) - .collect(Collectors.toList()); + .mapToObj(i -> generateTopicName(config.topicPrefix, config.partitionsPerTopic, i)) + .map(name -> new NewTopic(name, config.partitionsPerTopic, (short) 1).configs(config.topicConfigs)) + .collect(Collectors.toList()); + + int topicsPerBatch = MAX_PARTITIONS_PER_BATCH / config.partitionsPerTopic; + List> requests = Lists.partition(newTopics, topicsPerBatch); + + Map> results = requests.stream() + .map(admin::createTopics) + .map(CreateTopicsResult::values) + .flatMap(map -> map.entrySet().stream()) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + results.forEach(this::waitTopicCreated); + + return results.keySet().stream() + .map(name -> new Topic(name, config.partitionsPerTopic)) + .collect(Collectors.toList()); } /** - * Delete all topics except internal topics (starting with '__'). + * Delete all historical performance test topics. */ public int deleteTopics() { ListTopicsResult result = admin.listTopics(); try { - Set topics = result.names().get(); - topics.removeIf(name -> name.startsWith("__")); - admin.deleteTopics(topics).all().get(); - return topics.size(); + List toDelete = result.names().get().stream() + .filter(name -> name.startsWith(COMMON_TOPIC_PREFIX)) + .collect(Collectors.toList()); + admin.deleteTopics(toDelete).all().get(); + return toDelete.size(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } catch (ExecutionException ignored) { @@ -87,7 +120,9 @@ private void waitTopicCreated(String name, KafkaFuture future) { Thread.currentThread().interrupt(); } catch (ExecutionException e) { if (e.getCause() instanceof TopicExistsException) { - LOGGER.debug("Topic already exists. name={}", name); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("Topic already exists. name={}", name); + } } else { LOGGER.error("Failed to create topic. name={}", name, e); throw new RuntimeException(e); @@ -96,7 +131,58 @@ private void waitTopicCreated(String name, KafkaFuture future) { } private String generateTopicName(String topicPrefix, int partitions, int index) { - return String.format("%s-%04d-%07d", topicPrefix, partitions, index); + return String.format("%s%s_%04d_%07d", COMMON_TOPIC_PREFIX, topicPrefix, partitions, index); + } + + /** + * Find existing topics with the given prefix. + */ + public List findExistingTopicsByPrefix(String prefix) { + ListTopicsResult listTopicsResult = admin.listTopics(); + try { + String fullPrefix = COMMON_TOPIC_PREFIX + prefix; + + // Get all topic listings + Map topicListingsMap = listTopicsResult.namesToListings().get(); + + // Filter topic names by prefix + List matchingTopicNames = topicListingsMap.keySet().stream() + .filter(name -> name.startsWith(fullPrefix)) // Corrected fulalPrefix to fullPrefix + .collect(Collectors.toList()); + + if (matchingTopicNames.isEmpty()) { + return Collections.emptyList(); + } + + // Describe the filtered topics to get partition info + DescribeTopicsResult describeTopicsResult = admin.describeTopics(matchingTopicNames); + Map topicDescriptions = describeTopicsResult.allTopicNames().get(); + + return matchingTopicNames.stream() + .map(name -> { + TopicDescription description = topicDescriptions.get(name); + if (description != null) { + int partitionCount = description.partitions().size(); + return new Topic(name, partitionCount); + } else { + // This case should ideally not be reached if a topic name from listTopics is + // passed to describeTopics + // and the topic still exists. + LOGGER.warn("Could not find description for topic: {}. It might have been deleted.", name); + return null; + } + }) + .filter(Objects::nonNull) // Filter out any nulls if a description was missing + .collect(Collectors.toList()); + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + // It's good practice to specify which operation was interrupted. + throw new RuntimeException("Interrupted while listing or describing topics", e); + } catch (ExecutionException e) { + // Similarly, specify the operation. + throw new RuntimeException("Failed to list or describe topics", e); + } } public static class TopicsConfig { @@ -122,10 +208,18 @@ public Topic(String name, int partitions) { this.partitions = partitions; } + public String name() { + return name; + } + public List partitions() { return IntStream.range(0, partitions) - .mapToObj(i -> new TopicPartition(name, i)) - .collect(Collectors.toList()); + .mapToObj(i -> new TopicPartition(name, i)) + .collect(Collectors.toList()); + } + + public TopicPartition firstPartition() { + return new TopicPartition(name, 0); } public boolean containsPartition(TopicPartition partition) { diff --git a/tools/src/main/java/org/apache/kafka/tools/automq/perf/UniformRateLimiter.java b/tools/src/main/java/org/apache/kafka/tools/automq/perf/UniformRateLimiter.java index ef808663ec..085ad18799 100644 --- a/tools/src/main/java/org/apache/kafka/tools/automq/perf/UniformRateLimiter.java +++ b/tools/src/main/java/org/apache/kafka/tools/automq/perf/UniformRateLimiter.java @@ -1,12 +1,20 @@ /* - * Copyright 2024, AutoMQ HK Limited. + * Copyright 2025, AutoMQ HK Limited. * - * The use of this file is governed by the Business Source License, - * as detailed in the file "/LICENSE.S3Stream" included in this repository. + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * As of the Change Date specified in that file, in accordance with - * the Business Source License, use of this software will be governed - * by the Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.kafka.tools.automq.perf;